Materialize immutable-segment columns lazily behind an opt-in instance config (default off)

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

@@ -0,0 +1,125 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.segment.local.indexsegment.immutable;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.local.segment.index.column.PhysicalColumnIndexContainer;
+import org.apache.pinot.segment.local.segment.index.readers.text.MultiColumnLuceneTextIndexReader;
+import org.apache.pinot.segment.spi.ColumnMetadata;
+import org.apache.pinot.segment.spi.index.FieldIndexConfigs;
+import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer;
+import org.apache.pinot.segment.spi.store.SegmentDirectory;
+
+
+/// Creates the [ColumnIndexContainer] of a physical column of an [ImmutableSegmentImpl] on demand.
+///
+/// [ImmutableSegmentLoader] builds one per segment when lazy column materialization is on, instead of a
+/// [PhysicalColumnIndexContainer] per column at load. It retains only what creating a container later needs: the
+/// segment reader (held for the segment's lifetime anyway), the forward-index-only flag, the shared multi-column text
+/// index reader and, per column, the [FieldIndexConfigs] that were in effect at load.
+///
+/// The per-column configs are snapshotted at construction because the loading config they come from is mutable and
+/// shared: its map is replaced whenever the config is refreshed and mutated in place while OPEN_STRUCT child configs
+/// are resolved. The snapshot is compacted so that a wide segment retains close to nothing per column, which is the
+/// point of materializing lazily: configs that are equal by value collapse to one instance, the most common one
+/// becomes the implicit default, and only the columns that differ from it keep an entry (keyed by the column-name
+/// strings the segment metadata already holds). A column absent from the loading config maps to
+/// [FieldIndexConfigs#EMPTY], exactly what the eager path hands to the container. Collapsing relies on the value
+/// equality of the index configs; a config type that inherits the enabled/disabled-only equality of `IndexConfig`
+/// collapses on that alone, which is safe as long as its reader factory ignores the config (true of OPEN_STRUCT, the
+/// one such type today).
+///
+/// Thread-safe: immutable after construction, and creating a container mutates nothing here.
+class ColumnMaterializer {
+  private final SegmentDirectory.Reader _segmentReader;
+  private final boolean _forwardIndexOnly;
+  private final FieldIndexConfigs _defaultFieldIndexConfigs;
+  private final Map<String, FieldIndexConfigs> _fieldIndexConfigOverrides;
+  @Nullable
+  private final MultiColumnLuceneTextIndexReader _multiColumnTextIndex;
+  private final Set<String> _multiColumnTextIndexColumns;
+
+  /// @param columns the physical columns of the segment; their configs are looked up in `fieldIndexConfigByColumn`
+  ///                now, so the map may change afterwards
+  /// @param multiColumnTextIndexColumns the columns covered by `multiColumnTextIndex` (empty when there is none)
+  ColumnMaterializer(SegmentDirectory.Reader segmentReader, Collection<String> columns,
+      Map<String, FieldIndexConfigs> fieldIndexConfigByColumn, boolean forwardIndexOnly,
+      @Nullable MultiColumnLuceneTextIndexReader multiColumnTextIndex, Set<String> multiColumnTextIndexColumns) {
+    _segmentReader = segmentReader;
+    _forwardIndexOnly = forwardIndexOnly;
+    _multiColumnTextIndex = multiColumnTextIndex;
+    _multiColumnTextIndexColumns = multiColumnTextIndexColumns;
+
+    Map<FieldIndexConfigs, FieldIndexConfigs> canonical = new HashMap<>();
+    Map<FieldIndexConfigs, Integer> counts = new HashMap<>();
+    Map<String, FieldIndexConfigs> configsByColumn = new HashMap<>();
+    for (String column : columns) {
+      FieldIndexConfigs configs = canonical.computeIfAbsent(
+          fieldIndexConfigByColumn.getOrDefault(column, FieldIndexConfigs.EMPTY), Function.identity());
+      counts.merge(configs, 1, Integer::sum);
+      configsByColumn.put(column, configs);
+    }
+    FieldIndexConfigs defaultConfigs = FieldIndexConfigs.EMPTY;
+    int maxCount = 0;
+    for (Map.Entry<FieldIndexConfigs, Integer> entry : counts.entrySet()) {
+      if (entry.getValue() > maxCount) {
+        defaultConfigs = entry.getKey();
+        maxCount = entry.getValue();
+      }
+    }
+    _defaultFieldIndexConfigs = defaultConfigs;
+    configsByColumn.values().removeIf(configs -> configs == _defaultFieldIndexConfigs);
+    _fieldIndexConfigOverrides = Map.copyOf(configsByColumn);
+  }
+
+  /// Creates the index container of the column, attaching the shared multi-column text index reader when the column is
+  /// part of it. Fails with an [UncheckedIOException] when an index cannot be read.
+  ColumnIndexContainer createIndexContainer(ColumnMetadata columnMetadata) {
+    String column = columnMetadata.getColumnName();
+    PhysicalColumnIndexContainer container;
+    try {
+      container = new PhysicalColumnIndexContainer(_segmentReader, columnMetadata, getFieldIndexConfigs(column),
+          _forwardIndexOnly);
+    } catch (IOException e) {
+      throw new UncheckedIOException("Failed to materialize the indexes of column: " + column, e);
+    }
+    if (_multiColumnTextIndex != null && _multiColumnTextIndexColumns.contains(column)) {
+      container.setMultiColumnTextIndex(_multiColumnTextIndex);
+    }
+    return container;
+  }
+
+  @VisibleForTesting
+  FieldIndexConfigs getFieldIndexConfigs(String column) {
+    return _fieldIndexConfigOverrides.getOrDefault(column, _defaultFieldIndexConfigs);
+  }
+
+  @VisibleForTesting
+  Map<String, FieldIndexConfigs> getFieldIndexConfigOverrides() {
+    return _fieldIndexConfigOverrides;
+  }
+}

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

@@ -24,12 +24,18 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import javax.annotation.Nullable;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.segment.local.dedup.PartitionDedupMetadataManager;
@@ -75,6 +81,27 @@
 import org.slf4j.LoggerFactory;
 
 
+/// Immutable segment served from a [SegmentDirectory].
+///
+/// Physical columns are materialized in one of two modes, chosen at load by
+/// [IndexLoadingConfig#isLazyColumnMaterialization()]:
+///
+/// - **Eager** (the default and the public constructors): the [ColumnIndexContainer] and [DataSource] of every column
+///   exist from construction, so an unreadable index fails the segment load.
+/// - **Lazy** (the package-private constructor, built by [ImmutableSegmentLoader]): a physical column's container and
+///   data source are created on its first access through [#getDataSourceNullable(String)] or
+///   [#getIndex(String, IndexType)], exactly once per column even under concurrent first access; a failed creation
+///   leaves no mapping behind and propagates to the caller, so an unreadable index of a never-queried column is
+///   reported on first access instead of at load. Built-in virtual columns, star-tree dimensions and the segment-level
+///   star-tree and multi-column text indexes keep their eager path, and OPEN_STRUCT child columns are materialized
+///   together with their parent. Whole-segment consumers (`SELECT *`, segment metadata and index listings, record
+///   readers without a projection) materialize every column of the segments they touch, which is never worse than the
+///   eager mode. [#destroy()] closes only what was materialized and refuses any later materialization.
+///
+/// Thread safety: query-path lookups are lock-free in both modes. Lazy materialization holds a read lock while it
+/// creates and registers a container, and [#destroy()] takes the write lock before it closes anything, so every
+/// container whose creation was in flight is registered and closed; callers that hold a reference through the segment
+/// data manager never race destroy() at all.
 public class ImmutableSegmentImpl implements ImmutableSegment {
   private static final Logger LOGGER = LoggerFactory.getLogger(ImmutableSegmentImpl.class);
 
@@ -84,6 +111,18 @@ public class ImmutableSegmentImpl implements ImmutableSegment {
   private final StarTreeIndexContainer _starTreeIndexContainer;
   private final TextIndexReader _multiColumnTextIndex;
   private final Map<String, DataSource> _dataSources;
+
+  // Lazy column materialization; all null in eager mode. See the class documentation.
+  @Nullable
+  private final ColumnMaterializer _columnMaterializer;
+  // OPEN_STRUCT parent column -> its materialized child columns (col$key, col$__sparse__), restricted to parents that
+  // the segment schema declares as complex; null when the segment has none.
+  @Nullable
+  private final Map<String, List<String>> _openStructChildren;
+  @Nullable
+  private final ReadWriteLock _materializationLock;
+  // Guarded by _materializationLock
+  private boolean _destroyed;
   // Guards the post-registration hook so it reaches the directory at most once per segment instance, even when the
   // same segment is registered more than once (e.g. an upsert replacement with a consistency mode other than NONE
   // registers the new segment through a DuoSegmentDataManager and then directly).
@@ -108,6 +147,9 @@ public ImmutableSegmentImpl(
     _segmentMetadata = segmentMetadata;
     _indexContainerMap = columnIndexContainerMap;
     _starTreeIndexContainer = starTreeIndexContainer;
+    _columnMaterializer = null;
+    _openStructChildren = null;
+    _materializationLock = null;
     _dataSources =
         new Object2ObjectOpenHashMap<>(segmentMetadata.getColumnMetadataMap().size());
 
@@ -166,6 +208,117 @@ public ImmutableSegmentImpl(
     this(segmentDirectory, segmentMetadata, columnIndexContainerMap, starTreeIndexContainer, null);
   }
 
+  /// Creates a segment that materializes its physical columns lazily through `columnMaterializer`.
+  ///
+  /// `materializedIndexContainers` holds the containers created at load (built-in virtual columns and star-tree
+  /// dimensions) and becomes the registry of every container created afterwards, so that [#destroy()] closes exactly
+  /// the materialized ones. The columns already in it get their data source now, as in the eager mode.
+  ImmutableSegmentImpl(SegmentDirectory segmentDirectory, SegmentMetadataImpl segmentMetadata,
+      ColumnMaterializer columnMaterializer, ConcurrentMap<String, ColumnIndexContainer> materializedIndexContainers,
+      @Nullable StarTreeIndexContainer starTreeIndexContainer,
+      @Nullable MultiColumnLuceneTextIndexReader multiColumnTextIndex) {
+    _segmentDirectory = segmentDirectory;
+    _segmentMetadata = segmentMetadata;
+    _indexContainerMap = materializedIndexContainers;
+    _starTreeIndexContainer = starTreeIndexContainer;
+    _multiColumnTextIndex = multiColumnTextIndex;
+    _columnMaterializer = columnMaterializer;
+    _openStructChildren = groupOpenStructChildren(segmentMetadata);
+    _materializationLock = new ReentrantReadWriteLock();
+    _dataSources = new ConcurrentHashMap<>();
+    for (String column : materializedIndexContainers.keySet()) {
+      materializeDataSource(column);
+    }
+  }
+
+  /// Groups the materialized OPEN_STRUCT child columns under their parent, keeping only the parents the segment schema
+  /// declares as complex (the same rule the eager constructor applies).
+  @Nullable
+  private static Map<String, List<String>> groupOpenStructChildren(SegmentMetadataImpl segmentMetadata) {
+    Map<String, List<String>> children = null;
+    for (Map.Entry<String, ColumnMetadata> entry : segmentMetadata.getColumnMetadataMap().entrySet()) {
+      if (entry.getValue() instanceof ColumnMetadataImpl impl && impl.isMaterializedChild()) {
+        if (children == null) {
+          children = new HashMap<>();
+        }
+        children.computeIfAbsent(impl.getParentColumn(), k -> new ArrayList<>()).add(entry.getKey());
+      }
+    }
+    if (children == null) {
+      return null;
+    }
+    Schema schema = segmentMetadata.getSchema();
+    children.keySet()
+        .removeIf(parent -> !(schema != null && schema.getFieldSpecFor(parent) instanceof ComplexFieldSpec));
+    return children.isEmpty() ? null : children;
+  }
+
+  /// Lazy mode: returns the data source of the column, creating it on first access, or `null` when the segment has no
+  /// such column. OPEN_STRUCT child columns are reachable only through their parent, as in the eager mode.
+  @Nullable
+  private DataSource materializeDataSource(String column) {
+    ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataMap().get(column);
+    boolean openStructParent = _openStructChildren != null && _openStructChildren.containsKey(column);
+    if (!openStructParent && (columnMetadata == null || isMaterializedChild(columnMetadata))) {
+      return null;
+    }
+    Lock lock = _materializationLock.readLock();
+    lock.lock();
+    try {
+      checkNotDestroyed(column);
+      // Single flight per column: the mapping function runs at most once per column and leaves no mapping when it
+      // fails. It never reads this map again (creating the children of an OPEN_STRUCT parent goes through
+      // _indexContainerMap only), which computeIfAbsent forbids.
+      return _dataSources.computeIfAbsent(column,
+          k -> openStructParent ? createOpenStructDataSource(k) : createDataSource(k, columnMetadata));
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  private DataSource createDataSource(String column, ColumnMetadata columnMetadata) {
+    ColumnIndexContainer container = materializedIndexContainer(column, columnMetadata);
+    return columnMetadata.getFieldSpec().getDataType() == FieldSpec.DataType.MAP
+        ? new ImmutableMapDataSource(columnMetadata, container) : new ImmutableDataSource(columnMetadata, container);
+  }
+
+  private DataSource createOpenStructDataSource(String parent) {
+    Map<String, ColumnMetadata> columnMetadataMap = _segmentMetadata.getColumnMetadataMap();
+    Map<String, DataSource> denseChildren = new HashMap<>();
+    DataSource sparseChild = null;
+    for (String child : _openStructChildren.get(parent)) {
+      ColumnMetadata childMetadata = columnMetadataMap.get(child);
+      DataSource childDataSource =
+          new ImmutableDataSource(childMetadata, materializedIndexContainer(child, childMetadata));
+      if (OpenStructNaming.isSparseColumn(child)) {
+        sparseChild = childDataSource;
+      } else {
+        denseChildren.put(OpenStructNaming.parseKey(child), childDataSource);
+      }
+    }
+    ComplexFieldSpec fieldSpec = (ComplexFieldSpec) _segmentMetadata.getSchema().getFieldSpecFor(parent);
+    List<String> sparseKeys =
+        columnMetadataMap.get(parent) instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null;
+    return new ImmutableOpenStructDataSource(fieldSpec, denseChildren, sparseChild, _segmentMetadata.getTotalDocs(),
+        sparseKeys);
+  }
+
+  /// Lazy mode: returns the index container of the column, creating and registering it on first access. The mapping
+  /// function opens the column's index readers while it holds the map's bin lock, so a slow open (e.g. an on-heap
+  /// dictionary) can briefly stall the first access to an unrelated column in the same bin.
+  private ColumnIndexContainer materializedIndexContainer(String column, ColumnMetadata columnMetadata) {
+    return _indexContainerMap.computeIfAbsent(column, k -> _columnMaterializer.createIndexContainer(columnMetadata));
+  }
+
+  private void checkNotDestroyed(String column) {
+    Preconditions.checkState(!_destroyed, "Cannot materialize column: %s of destroyed segment: %s", column,
+        getSegmentName());
+  }
+
+  private static boolean isMaterializedChild(ColumnMetadata columnMetadata) {
+    return columnMetadata instanceof ColumnMetadataImpl impl && impl.isMaterializedChild();
+  }
+
   public void enableDedup(PartitionDedupMetadataManager partitionDedupMetadataManager) {
     _partitionDedupMetadataManager = partitionDedupMetadataManager;
   }
@@ -247,6 +400,19 @@ public boolean isReloadNeeded(IndexLoadingConfig indexLoadingConfig)
   @Override
   public <I extends IndexReader> I getIndex(String column, IndexType<?, I, ?> type) {
     ColumnIndexContainer container = _indexContainerMap.get(column);
+    if (container == null && _columnMaterializer != null) {
+      ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataMap().get(column);
+      if (columnMetadata != null) {
+        Lock lock = _materializationLock.readLock();
+        lock.lock();
+        try {
+          checkNotDestroyed(column);
+          container = materializedIndexContainer(column, columnMetadata);
+        } finally {
+          lock.unlock();
+        }
+      }
+    }
     if (container == null) {
       throw new NullPointerException("Invalid column: " + column);
     }
@@ -355,6 +521,17 @@ public void offload() {
   public void destroy() {
     String segmentName = getSegmentName();
     LOGGER.info("Trying to destroy segment : {}", segmentName);
+    if (_materializationLock != null) {
+      // Waits for in-flight materialization to register its containers, then refuses any further one, so the loop
+      // below closes exactly the materialized containers
+      Lock lock = _materializationLock.writeLock();
+      lock.lock();
+      try {
+        _destroyed = true;
+      } finally {
+        lock.unlock();
+      }
+    }
     if (_partitionUpsertMetadataManager != null) {
       _partitionUpsertMetadataManager.untrackSegmentForUpsertView(this);
     }
@@ -391,7 +568,11 @@ public void destroy() {
   @Nullable
   @Override
   public DataSource getDataSourceNullable(String column) {
-    return _dataSources.get(column);
+    DataSource dataSource = _dataSources.get(column);
+    if (dataSource == null && _columnMaterializer != null) {
+      dataSource = materializeDataSource(column);
+    }
+    return dataSource;
   }
 
   @Nullable