Prune segments from column metadata instead of materializing them

pinot-core/src/main/java/org/apache/pinot/core/query/pruner/ColumnValueSegmentPruner.java

@@ -88,10 +88,7 @@ boolean pruneSegmentWithPredicate(IndexSegment segment, Predicate predicate, Map
   private boolean pruneEqPredicate(IndexSegment segment, EqPredicate eqPredicate,
       Map<String, DataSource> dataSourceCache, ValueCache valueCache, QueryContext query) {
     String column = eqPredicate.getLhs().getIdentifier();
-    DataSource dataSource = segment instanceof ImmutableSegment ? segment.getDataSource(column, query.getSchema())
-        : dataSourceCache.computeIfAbsent(column, col -> segment.getDataSource(column, query.getSchema()));
-    assert dataSource != null;
-    DataSourceMetadata dataSourceMetadata = dataSource.getDataSourceMetadata();
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata(segment, column, dataSourceCache, query);
     ValueCache.CachedValue cachedValue = valueCache.get(eqPredicate, dataSourceMetadata.getDataType());
     // Check min/max value
     if (!checkMinMaxRange(dataSourceMetadata, cachedValue.getComparableValue())) {
@@ -120,10 +117,7 @@ private boolean pruneInPredicate(IndexSegment segment, InPredicate inPredicate,
       return false;
     }
     String column = inPredicate.getLhs().getIdentifier();
-    DataSource dataSource = segment instanceof ImmutableSegment ? segment.getDataSource(column, query.getSchema())
-        : dataSourceCache.computeIfAbsent(column, col -> segment.getDataSource(column, query.getSchema()));
-    assert dataSource != null;
-    DataSourceMetadata dataSourceMetadata = dataSource.getDataSourceMetadata();
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata(segment, column, dataSourceCache, query);
     List<ValueCache.CachedValue> cachedValues = valueCache.get(inPredicate, dataSourceMetadata.getDataType());
     // Check min/max value
     for (ValueCache.CachedValue value : cachedValues) {
@@ -140,10 +134,7 @@ private boolean pruneInPredicate(IndexSegment segment, InPredicate inPredicate,
   private boolean pruneRangePredicate(IndexSegment segment, RangePredicate rangePredicate,
       Map<String, DataSource> dataSourceCache, QueryContext query) {
     String column = rangePredicate.getLhs().getIdentifier();
-    DataSource dataSource = segment instanceof ImmutableSegment ? segment.getDataSource(column, query.getSchema())
-        : dataSourceCache.computeIfAbsent(column, col -> segment.getDataSource(column, query.getSchema()));
-    assert dataSource != null;
-    DataSourceMetadata dataSourceMetadata = dataSource.getDataSourceMetadata();
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata(segment, column, dataSourceCache, query);
 
     // Get lower/upper boundary value
     DataType dataType = dataSourceMetadata.getDataType();
@@ -222,4 +213,18 @@ private boolean checkMinMaxRange(DataSourceMetadata dataSourceMetadata, Comparab
     }
     return true;
   }
+
+  /// Pruning reads only the column's statistics, so an immutable segment answers from its column metadata rather
+  /// than materializing the column. A mutable segment keeps the per-segment data-source cache it had, where the
+  /// lookup is a map read and the metadata is not derivable without the data source.
+  private static DataSourceMetadata getDataSourceMetadata(IndexSegment segment, String column,
+      Map<String, DataSource> dataSourceCache, QueryContext query) {
+    if (segment instanceof ImmutableSegment) {
+      return segment.getDataSourceMetadata(column, query.getSchema());
+    }
+    DataSource dataSource = dataSourceCache.computeIfAbsent(column,
+        col -> segment.getDataSource(col, query.getSchema()));
+    assert dataSource != null;
+    return dataSource.getDataSourceMetadata();
+  }
 }

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

@@ -56,6 +56,7 @@
 import org.apache.pinot.segment.spi.FetchContext;
 import org.apache.pinot.segment.spi.ImmutableSegment;
 import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
 import org.apache.pinot.segment.spi.index.IndexReader;
 import org.apache.pinot.segment.spi.index.IndexType;
 import org.apache.pinot.segment.spi.index.StandardIndexes;
@@ -451,6 +452,20 @@ public SegmentMetadataImpl getSegmentMetadata() {
     return _segmentMetadata;
   }
 
+  /// Answers from the column metadata, so a caller that needs only the column's statistics does not materialize the
+  /// column. That matters most under lazy column materialization: segment pruning asks this for every segment the
+  /// server holds, and building an index container per segment there would put a reader — and, for an external
+  /// table, a Parquet footer parse — on the query thread for segments that are about to be pruned away.
+  ///
+  /// Falls back to the data source for a column the segment does not have, which is where the schema-driven default
+  /// and virtual columns are created.
+  @Override
+  public DataSourceMetadata getDataSourceMetadata(String column, Schema schema) {
+    ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(column);
+    return columnMetadata != null ? ImmutableDataSource.metadataOf(columnMetadata)
+        : getDataSource(column, schema).getDataSourceMetadata();
+  }
+
   @Override
   public DataSource getDataSource(String column, Schema schema) {
     DataSource dataSource = getDataSourceNullable(column);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/datasource/ImmutableDataSource.java

@@ -34,6 +34,12 @@ public ImmutableDataSource(ColumnMetadata columnMetadata, ColumnIndexContainer c
     super(new ImmutableDataSourceMetadata(columnMetadata), columnIndexContainer);
   }
 
+  /// The data-source metadata for a column, without a data source. It delegates to [ColumnMetadata] and holds no
+  /// index readers, so a caller that needs only the column's statistics can avoid materializing the column.
+  public static DataSourceMetadata metadataOf(ColumnMetadata columnMetadata) {
+    return new ImmutableDataSourceMetadata(columnMetadata);
+  }
+
   /// Exposes the segment's [ColumnMetadata] through the [DataSourceMetadata] view by delegating every accessor.
   /// Holding a single reference instead of copying the ten fields the view exposes keeps this object at one
   /// reference per column, which matters for wide segments where every loaded column retains one. The delegation

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java

@@ -38,6 +38,7 @@
 import org.apache.pinot.segment.local.segment.virtualcolumn.DocIdVirtualColumnProvider;
 import org.apache.pinot.segment.spi.ColumnMetadata;
 import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
 import org.apache.pinot.segment.spi.index.StandardIndexes;
 import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer;
 import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl;
@@ -48,6 +49,7 @@
 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.utils.CommonConstants.Segment.BuiltInVirtualColumn;
 import org.testng.annotations.Test;
 
@@ -161,6 +163,36 @@ public void testLazyModeCreatesNothingAtConstruction()
     verify(segmentDirectory).close();
   }
 
+  /// Segment pruning reads a column's statistics for every segment the server holds, to decide which segments can
+  /// match at all. Reaching those statistics must not materialize the column: under lazy materialization that would
+  /// build an index container — and for an external table parse a Parquet footer — on the query thread, for a segment
+  /// that is about to be pruned away.
+  @Test
+  public void testDataSourceMetadataDoesNotMaterializeTheColumn()
+      throws Exception {
+    ColumnMetadataImpl a = columnMetadata(intColumn("a"), null);
+    ColumnMetadataImpl b = columnMetadata(intColumn("b"), null);
+    ColumnMaterializer materializer = mock(ColumnMaterializer.class);
+    SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+    ImmutableSegmentImpl segment = lazySegment(segmentDirectory, materializer, a, b);
+
+    DataSourceMetadata metadata = segment.getDataSourceMetadata("a", mock(Schema.class));
+    assertNotNull(metadata);
+    assertEquals(metadata.getFieldSpec(), a.getFieldSpec());
+    assertEquals(metadata.getDataType(), a.getDataType());
+    assertEquals(metadata.getNumDocs(), a.getTotalDocs());
+    assertEquals(metadata.isSorted(), a.isSorted());
+    // The whole point: reading the statistics built nothing.
+    verifyNoInteractions(materializer);
+
+    // It agrees with what the materialized data source reports, and only THAT materializes.
+    assertEquals(segment.getDataSource("a", mock(Schema.class)).getDataSourceMetadata().getDataType(),
+        metadata.getDataType());
+    verify(materializer, times(1)).createIndexContainer(a);
+
+    segment.destroy();
+  }
+
   @Test
   public void testLazyModeMaterializesEachColumnOnceUnderConcurrentAccess()
       throws Exception {

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

@@ -23,6 +23,7 @@
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
 import org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap;
 import org.apache.pinot.segment.spi.index.reader.TextIndexReader;
 import org.apache.pinot.segment.spi.index.startree.StarTreeV2;
@@ -71,6 +72,18 @@ default DataSource getDataSource(String column) {
   /// asked column.
   DataSource getDataSource(String column, Schema schema);
 
+  /// The metadata of a column's data source, for callers that need only the column's statistics — its data type,
+  /// min/max values, partitioning — and never read its values.
+  ///
+  /// Segment pruning is the motivating caller: it reads min/max to decide whether a segment can match at all, for
+  /// every segment the server holds. Going through [#getDataSource(String, Schema)] to reach that metadata forces an
+  /// implementation that builds its columns lazily to construct the whole index container — every index reader for
+  /// the column — for a segment it is about to discard. An implementation that can answer from column metadata alone
+  /// should override this; the default keeps the existing behaviour.
+  default DataSourceMetadata getDataSourceMetadata(String column, Schema schema) {
+    return getDataSource(column, schema).getDataSourceMetadata();
+  }
+
   /// Returns a list of star-trees (V2), or null if there is no star-tree (V2) in the segment.
   @Nullable
   List<StarTreeV2> getStarTrees();