Fix compacted column statistics to restrict all stats to valid documents (#18041)

Author: Jackie-Jiang

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/CommitTimeCompactionIntegrationTest.java

@@ -332,7 +332,7 @@ public void testCommitTimeCompactionWithMultiValueColumns()
       throws Exception {
     // Test Case: Multi-Value Fields with Commit-Time Compaction
     // Goal: Ensure commit-time compaction correctly handles multi-value dictionary columns
-    // (like arrays/lists) during segment conversion, validating the fix for CompactedDictEncodedColumnStatistics
+    // (like arrays/lists) during segment conversion, validating the fix for CompactedColumnStatistics
 
     // Create test data with multi-value fields similar to user's "tags" column
     List<String> testRecords = List.of("200,Player200,game1,85.5,1681054200000,false,action;shooter",

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/converter/RealtimeSegmentConverter.java

@@ -26,7 +26,7 @@
 import org.apache.pinot.common.metrics.ServerMeter;
 import org.apache.pinot.common.metrics.ServerMetrics;
 import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImpl;
-import org.apache.pinot.segment.local.realtime.converter.stats.RealtimeSegmentSegmentCreationDataSource;
+import org.apache.pinot.segment.local.realtime.converter.stats.MutableSegmentCreationDataSource;
 import org.apache.pinot.segment.local.segment.creator.TransformPipeline;
 import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
 import org.apache.pinot.segment.local.segment.readers.CompactedPinotSegmentRecordReader;
@@ -42,6 +42,7 @@
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.RecordReader;
+import org.roaringbitmap.RoaringBitmap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -122,24 +123,22 @@ public void build(@Nullable SegmentVersion segmentVersion, @Nullable ServerMetri
 
     if (useCompactedReader) {
       // Take a snapshot of validDocIds at the beginning of conversion to ensure consistency
-      ThreadSafeMutableRoaringBitmap validDocIdsSnapshot = getValidDocIdSnapshot();
-      if (validDocIdsSnapshot == null) {
+      RoaringBitmap validDocIds = getValidDocIds();
+      if (validDocIds == null) {
         throw new IllegalStateException("Cannot use CompactedPinotSegmentRecordReader without valid document IDs. "
             + "Segment may be corrupted.");
       }
       // Use CompactedPinotSegmentRecordReader to remove obsolete/invalidated records
-      try (CompactedPinotSegmentRecordReader recordReader = new CompactedPinotSegmentRecordReader(
-          validDocIdsSnapshot)) {
+      try (CompactedPinotSegmentRecordReader recordReader = new CompactedPinotSegmentRecordReader(validDocIds)) {
         recordReader.init(_realtimeSegmentImpl, sortedDocIds);
-        buildSegmentWithReader(driver, genConfig, recordReader, sortedDocIds, sortedColumn, useCompactedReader,
-            validDocIdsSnapshot);
+        buildSegmentWithReader(driver, genConfig, recordReader, sortedDocIds, sortedColumn, validDocIds);
         publishCompactionMetrics(serverMetrics, preCommitRowCount, driver, compactionStartTime);
       }
     } else {
       // Use regular PinotSegmentRecordReader (existing behavior)
       try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) {
         recordReader.init(_realtimeSegmentImpl, sortedDocIds);
-        buildSegmentWithReader(driver, genConfig, recordReader, sortedDocIds, sortedColumn, useCompactedReader, null);
+        buildSegmentWithReader(driver, genConfig, recordReader, sortedDocIds, sortedColumn, null);
       }
     }
 
@@ -152,13 +151,10 @@ public void build(@Nullable SegmentVersion segmentVersion, @Nullable ServerMetri
     }
   }
 
-  private @Nullable ThreadSafeMutableRoaringBitmap getValidDocIdSnapshot() {
-    ThreadSafeMutableRoaringBitmap validDocIdsSnapshot = null;
-    if (_realtimeSegmentImpl.getValidDocIds() != null) {
-      validDocIdsSnapshot = new ThreadSafeMutableRoaringBitmap(
-          _realtimeSegmentImpl.getValidDocIds().getMutableRoaringBitmap());
-    }
-    return validDocIdsSnapshot;
+  @Nullable
+  private RoaringBitmap getValidDocIds() {
+    ThreadSafeMutableRoaringBitmap validDocIds = _realtimeSegmentImpl.getValidDocIds();
+    return validDocIds != null ? validDocIds.getMutableRoaringBitmap().toRoaringBitmap() : null;
   }
 
   /**
@@ -202,28 +198,17 @@ private void publishCompactionMetrics(@Nullable ServerMetrics serverMetrics,
    * Common method to build segment with the provided record reader
    */
   private void buildSegmentWithReader(SegmentIndexCreationDriverImpl driver, SegmentGeneratorConfig genConfig,
-      RecordReader recordReader, int[] sortedDocIds, @Nullable String sortedColumn, boolean useCompactedReader,
-      @Nullable ThreadSafeMutableRoaringBitmap validDocIdsSnapshot)
+      RecordReader recordReader, int[] sortedDocIds, @Nullable String sortedColumn, @Nullable RoaringBitmap validDocIds)
       throws Exception {
-    RealtimeSegmentSegmentCreationDataSource dataSource;
-    if (useCompactedReader) {
-      // For compacted readers, use the constructor that takes sortedDocIds and pass the validDocIds snapshot
-      dataSource = new RealtimeSegmentSegmentCreationDataSource(_realtimeSegmentImpl, recordReader, sortedDocIds,
-          sortedColumn, validDocIdsSnapshot);
-    } else {
-      // For regular readers, use the original constructor
-      dataSource = new RealtimeSegmentSegmentCreationDataSource(_realtimeSegmentImpl,
-          (PinotSegmentRecordReader) recordReader, sortedColumn);
-    }
-    // initializes reader
+    MutableSegmentCreationDataSource dataSource =
+        new MutableSegmentCreationDataSource(_realtimeSegmentImpl, recordReader, sortedDocIds, sortedColumn,
+            validDocIds);
     driver.init(genConfig, dataSource, TransformPipeline.getPassThroughPipeline(_tableName), InstanceType.SERVER);
 
     if (!_enableColumnMajor) {
       driver.build();
     } else {
-      //buildByColumn uses validDocIds to skip invalid record while indexing each column. We pass the validDocIds
-      // only if we are using compacted reader.
-      driver.buildByColumn(_realtimeSegmentImpl, useCompactedReader ? validDocIdsSnapshot : null);
+      driver.buildByColumn(_realtimeSegmentImpl, validDocIds);
     }
   }