diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ColumnMaterializer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ColumnMaterializer.java new file mode 100644 index 000000000000..c25bdeef39ca --- /dev/null +++ b/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 _fieldIndexConfigOverrides; + @Nullable + private final MultiColumnLuceneTextIndexReader _multiColumnTextIndex; + private final Set _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 columns, + Map fieldIndexConfigByColumn, boolean forwardIndexOnly, + @Nullable MultiColumnLuceneTextIndexReader multiColumnTextIndex, Set multiColumnTextIndexColumns) { + _segmentReader = segmentReader; + _forwardIndexOnly = forwardIndexOnly; + _multiColumnTextIndex = multiColumnTextIndex; + _multiColumnTextIndexColumns = multiColumnTextIndexColumns; + + Map canonical = new HashMap<>(); + Map counts = new HashMap<>(); + Map 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 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 getFieldIndexConfigOverrides() { + return _fieldIndexConfigOverrides; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java index 85d3a0c79a3a..b295cf68a3c3 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java +++ b/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 _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> _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 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> groupOpenStructChildren(SegmentMetadataImpl segmentMetadata) { + Map> children = null; + for (Map.Entry 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 columnMetadataMap = _segmentMetadata.getColumnMetadataMap(); + Map 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 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 getIndex(String column, IndexType 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 diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java index 724bd04ef693..e731fb82a561 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java @@ -21,9 +21,12 @@ import com.google.common.base.Preconditions; import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap; import java.io.File; +import java.io.IOException; import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import javax.annotation.Nullable; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.segment.local.segment.index.column.PhysicalColumnIndexContainer; @@ -233,6 +236,14 @@ public static ImmutableSegment load(SegmentDirectory segmentDirectory, IndexLoad } SegmentDirectory.Reader segmentReader = segmentDirectory.createReader(); + String segmentName = segmentMetadata.getName(); + if (indexLoadingConfig.isLazyColumnMaterialization()) { + ImmutableSegmentImpl segment = + loadWithLazyColumns(segmentDirectory, segmentReader, segmentMetadata, indexLoadingConfig); + LOGGER.info("Successfully loaded segment: {} with SegmentDirectory, materializing columns lazily", segmentName); + return segment; + } + Map indexContainerMap = new Object2ObjectOpenHashMap<>(columnMetadataMap.size()); for (Map.Entry entry : columnMetadataMap.entrySet()) { // FIXME: text-index only works with local SegmentDirectory @@ -240,20 +251,7 @@ public static ImmutableSegment load(SegmentDirectory segmentDirectory, IndexLoad new PhysicalColumnIndexContainer(segmentReader, entry.getValue(), indexLoadingConfig)); } - // Instantiate virtual columns - String segmentName = segmentMetadata.getName(); - Schema segmentSchema = segmentMetadata.getSchema(); - VirtualColumnProviderFactory.addBuiltInVirtualColumnsToSegmentSchema(segmentSchema, segmentName); - for (FieldSpec fieldSpec : segmentSchema.getAllFieldSpecs()) { - if (fieldSpec.isVirtualColumn()) { - String columnName = fieldSpec.getName(); - VirtualColumnContext context = - new VirtualColumnContext(fieldSpec, segmentMetadata.getTotalDocs(), segmentMetadata); - VirtualColumnProvider provider = VirtualColumnProviderFactory.buildProvider(context); - indexContainerMap.put(columnName, provider.buildColumnIndexContainer(context)); - columnMetadataMap.put(columnName, provider.buildMetadata(context)); - } - } + instantiateVirtualColumns(segmentMetadata, indexContainerMap); // Load star-tree index if it exists StarTreeIndexContainer starTreeIndexContainer = null; @@ -279,6 +277,58 @@ public static ImmutableSegment load(SegmentDirectory segmentDirectory, IndexLoad return segment; } + /// Lazy counterpart of the load above (see [ImmutableSegmentImpl]): no per-column container is created here. The + /// built-in virtual columns keep their eager containers, the star-tree dimensions are materialized now because the + /// star-tree shares their dictionaries, and every other physical column waits for its first access. The + /// [ColumnMaterializer] snapshots the per-column index configs before the virtual columns are added to the metadata, + /// so it covers exactly the physical columns. + private static ImmutableSegmentImpl loadWithLazyColumns(SegmentDirectory segmentDirectory, + SegmentDirectory.Reader segmentReader, SegmentMetadataImpl segmentMetadata, + IndexLoadingConfig indexLoadingConfig) + throws IOException { + Map columnMetadataMap = segmentMetadata.getColumnMetadataMap(); + MultiColumnLuceneTextIndexReader mcTextReader = null; + Set mcTextColumns = Set.of(); + if (segmentReader.hasMultiColumnTextIndex()) { + mcTextReader = new MultiColumnLuceneTextIndexReader(segmentMetadata); + mcTextColumns = Set.copyOf(segmentMetadata.getMultiColumnTextMetadata().getColumns()); + } + ColumnMaterializer columnMaterializer = new ColumnMaterializer(segmentReader, columnMetadataMap.keySet(), + indexLoadingConfig.getFieldIndexConfigByColName(), indexLoadingConfig.isForwardIndexOnly(), mcTextReader, + mcTextColumns); + + ConcurrentMap indexContainerMap = new ConcurrentHashMap<>(); + instantiateVirtualColumns(segmentMetadata, indexContainerMap); + + StarTreeIndexContainer starTreeIndexContainer = null; + if (segmentReader.hasStarTreeIndex()) { + starTreeIndexContainer = new StarTreeIndexContainer(segmentReader, segmentMetadata, + column -> indexContainerMap.computeIfAbsent(column, + k -> columnMaterializer.createIndexContainer(columnMetadataMap.get(k)))); + } + + return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata, columnMaterializer, indexContainerMap, + starTreeIndexContainer, mcTextReader); + } + + /// Adds the built-in virtual columns to the segment schema and creates their index containers and metadata. + private static void instantiateVirtualColumns(SegmentMetadataImpl segmentMetadata, + Map indexContainerMap) { + Map columnMetadataMap = segmentMetadata.getColumnMetadataMap(); + Schema segmentSchema = segmentMetadata.getSchema(); + VirtualColumnProviderFactory.addBuiltInVirtualColumnsToSegmentSchema(segmentSchema, segmentMetadata.getName()); + for (FieldSpec fieldSpec : segmentSchema.getAllFieldSpecs()) { + if (fieldSpec.isVirtualColumn()) { + String columnName = fieldSpec.getName(); + VirtualColumnContext context = + new VirtualColumnContext(fieldSpec, segmentMetadata.getTotalDocs(), segmentMetadata); + VirtualColumnProvider provider = VirtualColumnProviderFactory.buildProvider(context); + indexContainerMap.put(columnName, provider.buildColumnIndexContainer(context)); + columnMetadataMap.put(columnName, provider.buildMetadata(context)); + } + } + } + /// Check segment directory against the IndexLoadingConfig to see if any preprocessing is needed, such as changing /// segment format, adding new indices or updating default columns. public static boolean needPreprocess(SegmentDirectory segmentDirectory, IndexLoadingConfig indexLoadingConfig) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java index 9350ebe35185..e09749594297 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/PhysicalColumnIndexContainer.java @@ -75,9 +75,17 @@ public final class PhysicalColumnIndexContainer implements ColumnIndexContainer public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, ColumnMetadata metadata, IndexLoadingConfig indexLoadingConfig) throws IOException { + this(segmentReader, metadata, indexLoadingConfig.getFieldIndexConfig(metadata.getColumnName()), + indexLoadingConfig.isForwardIndexOnly()); + } + + /// Creates the container from the column's own index configs (`null` meaning none, i.e. every index type at its + /// default) and the forward-index-only flag, without a reference to the loading config they were taken from. + public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, ColumnMetadata metadata, + @Nullable FieldIndexConfigs fieldIndexConfigs, boolean forwardIndexOnly) + throws IOException { String columnName = metadata.getColumnName(); - FieldIndexConfigs fieldIndexConfigs = indexLoadingConfig.getFieldIndexConfig(columnName); if (fieldIndexConfigs == null) { fieldIndexConfigs = FieldIndexConfigs.EMPTY; } @@ -92,7 +100,6 @@ public PhysicalColumnIndexContainer(SegmentDirectory.Reader segmentReader, Colum // Scratch array indexed by numeric id; compacted into the exactly-sized _readers below. IndexReader[] readersById = new IndexReader[numIndexTypes]; long presentMask = 0L; - boolean forwardIndexOnly = indexLoadingConfig.isForwardIndexOnly(); try { for (IndexType indexType : allIndexes) { if (forwardIndexOnly && !FORWARD_INDEX_ONLY_TYPES.contains(indexType.getId())) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfig.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfig.java index 3782f3abbf8e..c7bd326e65d6 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfig.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfig.java @@ -79,6 +79,7 @@ public class IndexLoadingConfig { private String _segmentStoreURI; private String _segmentDirectoryLoader; private Map> _instanceTierConfigs; + private boolean _lazyColumnMaterialization; // Initialized by table config and schema private List _sortedColumns = List.of(); @@ -168,6 +169,7 @@ private void extractFromInstanceConfig() { Map> tierConfigs = _instanceDataManagerConfig.getTierConfigs(); _instanceTierConfigs = tierConfigs != null ? tierConfigs : Map.of(); + _lazyColumnMaterialization = _instanceDataManagerConfig.isLazyColumnMaterialization(); } private void extractFromTableConfigAndSchema() { @@ -351,6 +353,18 @@ public void setForwardIndexOnly(boolean forwardIndexOnly) { _forwardIndexOnly = forwardIndexOnly; } + /// Whether immutable segments materialize the index container and data source of a physical column on its first + /// access instead of for every column at load. Sourced from + /// [InstanceDataManagerConfig#isLazyColumnMaterialization()] (off by default), so a segment loaded without an + /// instance config is always eager unless a caller opts in through [#setLazyColumnMaterialization(boolean)]. + public boolean isLazyColumnMaterialization() { + return _lazyColumnMaterialization; + } + + public void setLazyColumnMaterialization(boolean lazyColumnMaterialization) { + _lazyColumnMaterialization = lazyColumnMaterialization; + } + public boolean isSkipSegmentPreprocess() { if (_dirty) { refreshIndexConfigs(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexContainer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexContainer.java index 7dd16d364ba8..1711cde33b2e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexContainer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexContainer.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.List; import java.util.Map; +import java.util.function.Function; import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.startree.StarTreeV2; @@ -35,7 +36,15 @@ public class StarTreeIndexContainer implements Closeable { public StarTreeIndexContainer(SegmentDirectory.Reader segmentReader, SegmentMetadataImpl segmentMetadata, Map indexContainerMap) throws IOException { - _starTrees = StarTreeLoaderUtils.loadStarTreeV2(segmentReader, segmentMetadata, indexContainerMap); + this(segmentReader, segmentMetadata, indexContainerMap::get); + } + + /// `indexContainerProvider` returns the index container of a star-tree dimension column, which lets a segment that + /// materializes columns lazily create the container of each dimension while the star-tree is loaded. + public StarTreeIndexContainer(SegmentDirectory.Reader segmentReader, SegmentMetadataImpl segmentMetadata, + Function indexContainerProvider) + throws IOException { + _starTrees = StarTreeLoaderUtils.loadStarTreeV2(segmentReader, segmentMetadata, indexContainerProvider); } public List getStarTrees() { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java index 4e6ab15b49bb..9b99a4f8c5a4 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java @@ -23,6 +23,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Function; import org.apache.pinot.segment.local.aggregator.ValueAggregatorFactory; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexReaderFactory; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedBitSVForwardIndexReaderV2; @@ -52,6 +53,14 @@ private StarTreeLoaderUtils() { public static List loadStarTreeV2(SegmentDirectory.Reader segmentReader, SegmentMetadataImpl segmentMetadata, Map indexContainerMap) throws IOException { + return loadStarTreeV2(segmentReader, segmentMetadata, indexContainerMap::get); + } + + /// `indexContainerProvider` returns the index container of a dimension column; it is applied once per dimension + /// of each star-tree to fetch the dictionary the star-tree data source shares with the column. + public static List loadStarTreeV2(SegmentDirectory.Reader segmentReader, + SegmentMetadataImpl segmentMetadata, Function indexContainerProvider) + throws IOException { List starTreeMetadataList = segmentMetadata.getStarTreeV2MetadataList(); assert starTreeMetadataList != null; int numStarTrees = starTreeMetadataList.size(); @@ -72,7 +81,7 @@ public static List loadStarTreeV2(SegmentDirectory.Reader segmentRea FixedBitSVForwardIndexReaderV2 forwardIndex = new FixedBitSVForwardIndexReaderV2(forwardIndexDataBuffer, numDocs, columnMetadata.getBitsPerElement()); dataSourceMap.put(dimension, new StarTreeDataSource(columnMetadata.getFieldSpec(), numDocs, forwardIndex, - indexContainerMap.get(dimension).getIndex(StandardIndexes.dictionary()))); + indexContainerProvider.apply(dimension).getIndex(StandardIndexes.dictionary()))); } // Load metric (function-column pair) forward indexes diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ColumnMaterializerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ColumnMaterializerTest.java new file mode 100644 index 000000000000..b770d11bcaf4 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ColumnMaterializerTest.java @@ -0,0 +1,124 @@ +/** + * 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 java.io.IOException; +import java.io.UncheckedIOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +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.index.DictionaryIndexConfig; +import org.apache.pinot.segment.spi.index.FieldIndexConfigs; +import org.apache.pinot.segment.spi.index.StandardIndexes; +import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl; +import org.apache.pinot.segment.spi.store.SegmentDirectory; +import org.apache.pinot.spi.config.table.IndexConfig; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class ColumnMaterializerTest { + + /// The snapshot must retain close to nothing per column: configs equal by value collapse to one instance, the most + /// common one is the implicit default, only the other columns keep an entry, keyed by the caller's own strings. + @Test + public void testFieldIndexConfigsAreCanonicalizedAndCompacted() { + FieldIndexConfigs common = + new FieldIndexConfigs.Builder().add(StandardIndexes.dictionary(), DictionaryIndexConfig.DEFAULT).build(); + FieldIndexConfigs commonCopy = + new FieldIndexConfigs.Builder().add(StandardIndexes.dictionary(), new DictionaryIndexConfig(false)).build(); + assertEquals(commonCopy, common); + assertNotSame(commonCopy, common); + FieldIndexConfigs inverted = + new FieldIndexConfigs.Builder(common).add(StandardIndexes.inverted(), IndexConfig.ENABLED).build(); + // Distinct string instance from the key the loading config holds: the segment metadata's own name must be kept + String c3 = new StringBuilder("c3").toString(); + Map liveConfigs = Map.of("c1", common, "c2", commonCopy, "c3", inverted); + + ColumnMaterializer materializer = new ColumnMaterializer(mock(SegmentDirectory.Reader.class), + List.of("c1", "c2", c3, "c4"), liveConfigs, false, null, Set.of()); + + assertSame(materializer.getFieldIndexConfigs("c1"), materializer.getFieldIndexConfigs("c2")); + assertEquals(materializer.getFieldIndexConfigs("c1"), common); + assertSame(materializer.getFieldIndexConfigs("c3"), inverted); + // Absent from the loading config means no config, exactly what the eager path uses + assertSame(materializer.getFieldIndexConfigs("c4"), FieldIndexConfigs.EMPTY); + Map overrides = materializer.getFieldIndexConfigOverrides(); + assertEquals(overrides.keySet(), Set.of("c3", "c4")); + for (String column : overrides.keySet()) { + if (column.equals("c3")) { + assertSame(column, c3); + } + } + } + + @Test + public void testNoColumnsMeansNoConfigs() { + ColumnMaterializer materializer = + new ColumnMaterializer(mock(SegmentDirectory.Reader.class), List.of(), Map.of(), false, null, Set.of()); + assertSame(materializer.getFieldIndexConfigs("any"), FieldIndexConfigs.EMPTY); + assertTrue(materializer.getFieldIndexConfigOverrides().isEmpty()); + } + + @Test + public void testCreateIndexContainerAttachesMultiColumnTextIndexToItsColumnsOnly() { + // A reader without any index yields an empty container, enough to observe the attachment + SegmentDirectory.Reader reader = mock(SegmentDirectory.Reader.class); + MultiColumnLuceneTextIndexReader multiColumnTextIndex = mock(MultiColumnLuceneTextIndexReader.class); + ColumnMaterializer materializer = + new ColumnMaterializer(reader, List.of("t", "u"), Map.of(), false, multiColumnTextIndex, Set.of("t")); + + PhysicalColumnIndexContainer t = (PhysicalColumnIndexContainer) materializer.createIndexContainer(metadata("t")); + assertSame(t.getMultiColumnTextIndex(), multiColumnTextIndex); + PhysicalColumnIndexContainer u = (PhysicalColumnIndexContainer) materializer.createIndexContainer(metadata("u")); + assertNull(u.getMultiColumnTextIndex()); + assertNull(t.getIndex(StandardIndexes.forward())); + } + + @Test + public void testCreateIndexContainerWrapsReadFailure() + throws IOException { + SegmentDirectory.Reader reader = mock(SegmentDirectory.Reader.class); + when(reader.hasIndexFor("a", StandardIndexes.forward())).thenReturn(true); + when(reader.getIndexFor("a", StandardIndexes.forward())).thenThrow(new IOException("disk")); + ColumnMaterializer materializer = new ColumnMaterializer(reader, List.of("a"), Map.of(), true, null, Set.of()); + + UncheckedIOException e = + expectThrows(UncheckedIOException.class, () -> materializer.createIndexContainer(metadata("a"))); + assertTrue(e.getMessage().contains("a"), e.getMessage()); + assertEquals(e.getCause().getMessage(), "disk"); + } + + private static ColumnMetadataImpl metadata(String column) { + return ColumnMetadataImpl.builder().setFieldSpec(new DimensionFieldSpec(column, FieldSpec.DataType.STRING, true)) + .setTotalDocs(10).setCardinality(10).setHasDictionary(false).build(); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java index db1ce1aafcbd..864615fdf77e 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java @@ -18,21 +18,60 @@ */ package org.apache.pinot.segment.local.indexsegment.immutable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.List; import java.util.Map; +import java.util.Set; import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.apache.pinot.segment.local.segment.index.map.ImmutableMapDataSource; +import org.apache.pinot.segment.local.segment.index.openstruct.ImmutableOpenStructDataSource; +import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +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; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; import org.apache.pinot.segment.spi.store.SegmentDirectory; +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.testng.annotations.Test; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; -/// Tests for the post-registration lifecycle hook of [ImmutableSegmentImpl]. +/// Tests for [ImmutableSegmentImpl]: the post-registration lifecycle hook, and the lazy column materialization mode +/// against a mocked [ColumnMaterializer]. public class ImmutableSegmentImplTest { + private static final int NUM_DOCS = 10; /// The hook must reach the directory at most once per segment instance: the same segment can be registered more than /// once (e.g. an upsert replacement with a consistency mode other than NONE registers it through a @@ -86,4 +125,287 @@ private static ImmutableSegmentImpl createSegment(SegmentDirectory segmentDirect when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>()); return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata, Map.of(), null); } + + /// The eager (flag off) mode is untouched: every data source exists from construction over the given containers. + @Test + public void testEagerModeCreatesDataSourcesAtConstruction() { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + ImmutableSegmentImpl segment = + new ImmutableSegmentImpl(mock(SegmentDirectory.class), segmentMetadata(schema(a), a), Map.of("a", containerA), + null); + + assertSame(segment.getDataSourceNullable("a").getIndexContainer(), containerA); + assertNull(segment.getDataSourceNullable("unknown")); + } + + @Test + public void testLazyModeCreatesNothingAtConstruction() + 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, schema(a, b), materializer, a, b); + + verifyNoInteractions(materializer); + // Column listings come from the metadata schema and never materialize anything + assertEquals(segment.getColumnNames(), Set.of("a", "b")); + assertEquals(segment.getPhysicalColumnNames(), Set.of("a", "b")); + // Neither does asking for a column the segment does not have + assertNull(segment.getDataSourceNullable("unknown")); + assertThrows(NullPointerException.class, () -> segment.getIndex("unknown", StandardIndexes.forward())); + verifyNoInteractions(materializer); + + segment.destroy(); + verify(segmentDirectory).close(); + } + + @Test + public void testLazyModeMaterializesEachColumnOnceUnderConcurrentAccess() + throws Exception { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + AtomicInteger creations = new AtomicInteger(); + when(materializer.createIndexContainer(a)).thenAnswer(invocation -> { + creations.incrementAndGet(); + // Widen the window in which every other caller must wait for this creation instead of starting its own + Thread.sleep(50); + return containerA; + }); + ImmutableSegmentImpl segment = lazySegment(mock(SegmentDirectory.class), schema(a), materializer, a); + + int numCallers = 16; + ExecutorService executor = Executors.newFixedThreadPool(numCallers); + DataSource first; + try { + CountDownLatch start = new CountDownLatch(1); + List> futures = new ArrayList<>(numCallers); + for (int i = 0; i < numCallers; i++) { + futures.add(executor.submit(() -> { + start.await(); + return segment.getDataSourceNullable("a"); + })); + } + start.countDown(); + first = futures.get(0).get(); + for (Future future : futures) { + assertSame(future.get(), first); + } + } finally { + executor.shutdownNow(); + } + + assertNotNull(first); + assertSame(first.getIndexContainer(), containerA); + assertEquals(creations.get(), 1); + verify(materializer, times(1)).createIndexContainer(a); + } + + @Test + public void testDestroyClosesOnlyMaterializedContainersAndRefusesLaterMaterialization() + throws Exception { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ColumnMetadataImpl b = columnMetadata(intColumn("b"), null); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + ColumnIndexContainer containerB = mock(ColumnIndexContainer.class); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + when(materializer.createIndexContainer(a)).thenReturn(containerA); + when(materializer.createIndexContainer(b)).thenReturn(containerB); + SegmentDirectory segmentDirectory = mock(SegmentDirectory.class); + ImmutableSegmentImpl segment = lazySegment(segmentDirectory, schema(a, b), materializer, a, b); + DataSource dataSourceA = segment.getDataSourceNullable("a"); + assertNotNull(dataSourceA); + + segment.destroy(); + + verify(containerA).close(); + verify(containerB, never()).close(); + verify(materializer, never()).createIndexContainer(b); + verify(segmentDirectory).close(); + // A column that was never touched can no longer be materialized ... + assertThrows(IllegalStateException.class, () -> segment.getDataSourceNullable("b")); + assertThrows(IllegalStateException.class, () -> segment.getIndex("b", StandardIndexes.forward())); + // ... while a materialized one keeps returning its (closed) data source as in the eager mode, and a column the + // segment does not have is still simply absent + assertSame(segment.getDataSourceNullable("a"), dataSourceA); + assertNull(segment.getDataSourceNullable("unknown")); + verify(materializer, times(1)).createIndexContainer(any()); + } + + @Test + public void testGetIndexMaterializesAndSharesTheContainerWithTheDataSource() { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ForwardIndexReader forwardIndex = mock(ForwardIndexReader.class); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + doReturn(forwardIndex).when(containerA).getIndex(StandardIndexes.forward()); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + when(materializer.createIndexContainer(a)).thenReturn(containerA); + ImmutableSegmentImpl segment = lazySegment(mock(SegmentDirectory.class), schema(a), materializer, a); + + assertSame(segment.getForwardIndex("a"), forwardIndex); + verify(materializer, times(1)).createIndexContainer(a); + + DataSource dataSource = segment.getDataSourceNullable("a"); + assertNotNull(dataSource); + assertSame(dataSource.getIndexContainer(), containerA); + assertSame(dataSource.getForwardIndex(), forwardIndex); + verify(materializer, times(1)).createIndexContainer(a); + } + + @Test + public void testLazyModeGroupsOpenStructChildrenUnderTheirParent() + throws Exception { + ComplexFieldSpec metrics = new ComplexFieldSpec("metrics", FieldSpec.DataType.OPEN_STRUCT, true, + Map.of("views", new DimensionFieldSpec("views", FieldSpec.DataType.LONG, true))); + String viewsColumn = OpenStructNaming.materializedColumnName("metrics", "views"); + String sparseColumn = OpenStructNaming.sparseColumnName("metrics"); + ColumnMetadataImpl parent = columnMetadata(metrics, null); + ColumnMetadataImpl views = columnMetadata(new DimensionFieldSpec(viewsColumn, FieldSpec.DataType.LONG, true), + "metrics"); + ColumnMetadataImpl sparse = columnMetadata(new DimensionFieldSpec(sparseColumn, FieldSpec.DataType.STRING, true), + "metrics"); + ColumnMetadataImpl dim = columnMetadata(intColumn("dim"), null); + ForwardIndexReader viewsForwardIndex = mock(ForwardIndexReader.class); + ColumnIndexContainer viewsContainer = mock(ColumnIndexContainer.class); + doReturn(viewsForwardIndex).when(viewsContainer).getIndex(StandardIndexes.forward()); + ColumnIndexContainer sparseContainer = mock(ColumnIndexContainer.class); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + when(materializer.createIndexContainer(views)).thenReturn(viewsContainer); + when(materializer.createIndexContainer(sparse)).thenReturn(sparseContainer); + // Only the parent and the regular column are in the schema; the children live in the segment metadata alone + ImmutableSegmentImpl segment = lazySegment(mock(SegmentDirectory.class), schema(metrics, dim.getFieldSpec()), + materializer, parent, views, sparse, dim); + + DataSource parentDataSource = segment.getDataSourceNullable("metrics"); + assertTrue(parentDataSource instanceof ImmutableOpenStructDataSource); + ImmutableOpenStructDataSource openStruct = (ImmutableOpenStructDataSource) parentDataSource; + assertTrue(openStruct.isMaterialized("views")); + assertSame(openStruct.getDataSource("views").getIndexContainer(), viewsContainer); + assertFalse(openStruct.isFullyMaterialized()); + assertSame(segment.getDataSourceNullable("metrics"), parentDataSource); + // Children are reachable only through their parent, as in the eager mode + assertNull(segment.getDataSourceNullable(viewsColumn)); + assertNull(segment.getDataSourceNullable(sparseColumn)); + // Both children were materialized exactly once, together with the parent, and getIndex reuses their containers + verify(materializer, times(1)).createIndexContainer(views); + verify(materializer, times(1)).createIndexContainer(sparse); + verify(materializer, never()).createIndexContainer(parent); + verify(materializer, never()).createIndexContainer(dim); + assertSame(segment.getIndex(viewsColumn, StandardIndexes.forward()), viewsForwardIndex); + verify(materializer, times(1)).createIndexContainer(views); + + segment.destroy(); + verify(viewsContainer).close(); + verify(sparseContainer).close(); + } + + @Test + public void testLazyModeMapColumnYieldsMapDataSource() { + ComplexFieldSpec mapSpec = new ComplexFieldSpec("m", FieldSpec.DataType.MAP, true, + Map.of(ComplexFieldSpec.KEY_FIELD, new DimensionFieldSpec("key", FieldSpec.DataType.STRING, true), + ComplexFieldSpec.VALUE_FIELD, new DimensionFieldSpec("value", FieldSpec.DataType.INT, true))); + ColumnMetadataImpl m = columnMetadata(mapSpec, null); + ColumnIndexContainer container = mock(ColumnIndexContainer.class); + doReturn(mock(ForwardIndexReader.class)).when(container).getIndex(StandardIndexes.forward()); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + when(materializer.createIndexContainer(m)).thenReturn(container); + ImmutableSegmentImpl segment = lazySegment(mock(SegmentDirectory.class), schema(mapSpec), materializer, m); + + DataSource dataSource = segment.getDataSourceNullable("m"); + assertTrue(dataSource instanceof ImmutableMapDataSource); + assertSame(dataSource.getIndexContainer(), container); + } + + @Test + public void testLazyModeFailedMaterializationLeavesNoMappingAndRetrySucceeds() + throws Exception { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + when(materializer.createIndexContainer(a)).thenThrow(new UncheckedIOException(new IOException("boom"))) + .thenReturn(containerA); + ImmutableSegmentImpl segment = lazySegment(mock(SegmentDirectory.class), schema(a), materializer, a); + + assertThrows(UncheckedIOException.class, () -> segment.getDataSourceNullable("a")); + DataSource dataSource = segment.getDataSourceNullable("a"); + assertNotNull(dataSource); + assertSame(dataSource.getIndexContainer(), containerA); + verify(materializer, times(2)).createIndexContainer(a); + + segment.destroy(); + verify(containerA, times(1)).close(); + } + + /// Containers created at load (virtual columns, star-tree dimensions) are handed in already materialized: their data + /// sources exist from construction, the materializer is never asked for them, and destroy() closes them. + @Test + public void testLazyModeMaterializedContainersGetDataSourcesAtConstruction() + throws Exception { + ColumnMetadataImpl a = columnMetadata(intColumn("a"), null); + ColumnMetadataImpl b = columnMetadata(intColumn("b"), null); + ColumnIndexContainer containerA = mock(ColumnIndexContainer.class); + ColumnMaterializer materializer = mock(ColumnMaterializer.class); + ConcurrentMap materialized = new ConcurrentHashMap<>(Map.of("a", containerA)); + ImmutableSegmentImpl segment = + new ImmutableSegmentImpl(mock(SegmentDirectory.class), segmentMetadata(schema(a, b), a, b), materializer, + materialized, null, null); + + verifyNoInteractions(materializer); + assertSame(segment.getDataSourceNullable("a").getIndexContainer(), containerA); + verifyNoInteractions(materializer); + + segment.destroy(); + verify(containerA).close(); + } + + private static ImmutableSegmentImpl lazySegment(SegmentDirectory segmentDirectory, Schema schema, + ColumnMaterializer materializer, ColumnMetadata... columns) { + return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata(schema, columns), materializer, + new ConcurrentHashMap<>(), null, null); + } + + private static SegmentMetadataImpl segmentMetadata(Schema schema, ColumnMetadata... columns) { + SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); + when(segmentMetadata.getName()).thenReturn("seg"); + when(segmentMetadata.getTotalDocs()).thenReturn(NUM_DOCS); + when(segmentMetadata.getSchema()).thenReturn(schema); + TreeMap columnMetadataMap = new TreeMap<>(); + for (ColumnMetadata column : columns) { + columnMetadataMap.put(column.getColumnName(), column); + } + when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + return segmentMetadata; + } + + private static Schema schema(FieldSpec... fieldSpecs) { + Schema schema = new Schema(); + for (FieldSpec fieldSpec : fieldSpecs) { + schema.addField(fieldSpec); + } + return schema; + } + + private static Schema schema(ColumnMetadata... columns) { + Schema schema = new Schema(); + for (ColumnMetadata column : columns) { + schema.addField(column.getFieldSpec()); + } + return schema; + } + + private static FieldSpec intColumn(String name) { + return new DimensionFieldSpec(name, FieldSpec.DataType.INT, true); + } + + private static ColumnMetadataImpl columnMetadata(FieldSpec fieldSpec, @Nullable String parentColumn) { + ColumnMetadataImpl.Builder builder = + ColumnMetadataImpl.builder().setFieldSpec(fieldSpec).setTotalDocs(NUM_DOCS).setCardinality(NUM_DOCS) + .setHasDictionary(false); + if (parentColumn != null) { + builder.setParentColumn(parentColumn); + } + return builder.build(); + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfigTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfigTest.java index 4edb1c383cfb..d146833d2315 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfigTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfigTest.java @@ -46,6 +46,20 @@ public class IndexLoadingConfigTest { private static final String TABLE_NAME = "table01"; + @Test + public void testLazyColumnMaterializationIsSourcedFromInstanceConfig() { + assertFalse(new IndexLoadingConfig().isLazyColumnMaterialization()); + + InstanceDataManagerConfig idmCfg = mock(InstanceDataManagerConfig.class); + when(idmCfg.getConfig()).thenReturn(new PinotConfiguration()); + when(idmCfg.isLazyColumnMaterialization()).thenReturn(true); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build(); + IndexLoadingConfig ilc = new IndexLoadingConfig(idmCfg, tableConfig, null); + assertTrue(ilc.isLazyColumnMaterialization()); + ilc.setLazyColumnMaterialization(false); + assertFalse(ilc.isLazyColumnMaterialization()); + } + @Test public void testCalculateIndexConfigsWithoutTierOverwrites() throws IOException { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/LoaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/LoaderTest.java index ac22a9a3e56c..a1832ae461e2 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/LoaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/LoaderTest.java @@ -22,12 +22,14 @@ import java.net.URL; import java.util.List; import java.util.Map; +import java.util.Objects; import javax.annotation.Nullable; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; 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.PinotSegmentColumnReader; import org.apache.pinot.segment.local.segment.store.SegmentLocalFSDirectory; import org.apache.pinot.segment.local.utils.SegmentOperationsThrottler; import org.apache.pinot.segment.local.utils.SegmentOperationsThrottlerSet; @@ -38,12 +40,17 @@ import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; +import org.apache.pinot.segment.spi.index.IndexService; +import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; +import org.apache.pinot.segment.spi.index.startree.StarTreeV2; import org.apache.pinot.segment.spi.store.SegmentDirectory; import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.DimensionFieldSpec; @@ -996,6 +1003,73 @@ public void testVectorIndexLoad() indexSegment.destroy(); } + /// Lazy column materialization must be observationally identical to the eager mode on a real segment: same columns, + /// same index presence, same values, and star-tree dimensions materialized at load sharing their dictionary with the + /// column's own data source. + @Test + public void testLazyColumnMaterializationParity() + throws Exception { + constructV1Segment(); + Schema schema = createSchema(); + List starTreeDimensions = List.of("column1", "column5"); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setSegmentVersion("v3").setInvertedIndexColumns(List.of("column1")).setStarTreeIndexConfigs( + List.of(new StarTreeIndexConfig(starTreeDimensions, null, List.of("COUNT__*"), null, 100))).build(); + tableConfig.getIndexingConfig().setEnableDynamicStarTreeCreation(true); + + // The eager load converts the segment to v3 and builds the star-tree and the inverted index; the lazy load of the + // same directory then finds nothing pending + ImmutableSegment eager = ImmutableSegmentLoader.load(_indexDir, new IndexLoadingConfig(tableConfig, schema), + SEGMENT_OPERATIONS_THROTTLER); + IndexLoadingConfig lazyIndexLoadingConfig = new IndexLoadingConfig(tableConfig, schema); + lazyIndexLoadingConfig.setLazyColumnMaterialization(true); + ImmutableSegment lazy = ImmutableSegmentLoader.load(_indexDir, lazyIndexLoadingConfig, + SEGMENT_OPERATIONS_THROTTLER); + try { + assertEquals(lazy.getSegmentMetadata().getVersion(), SegmentVersion.v3); + assertEquals(lazy.getColumnNames(), eager.getColumnNames()); + assertEquals(lazy.getPhysicalColumnNames(), eager.getPhysicalColumnNames()); + + List lazyStarTrees = lazy.getStarTrees(); + assertNotNull(lazyStarTrees); + assertEquals(lazyStarTrees.size(), 1); + assertEquals(eager.getStarTrees().size(), 1); + for (String dimension : starTreeDimensions) { + assertSame(lazyStarTrees.get(0).getDataSource(dimension).getDictionary(), lazy.getDictionary(dimension)); + assertSame(eager.getStarTrees().get(0).getDataSource(dimension).getDictionary(), + eager.getDictionary(dimension)); + } + + List> indexTypes = IndexService.getInstance().getAllIndexes(); + for (String column : eager.getColumnNames()) { + DataSource eagerDataSource = eager.getDataSource(column); + DataSource lazyDataSource = lazy.getDataSource(column); + DataSourceMetadata eagerMetadata = eagerDataSource.getDataSourceMetadata(); + DataSourceMetadata lazyMetadata = lazyDataSource.getDataSourceMetadata(); + assertEquals(lazyMetadata.getFieldSpec(), eagerMetadata.getFieldSpec(), column); + assertEquals(lazyMetadata.getNumDocs(), eagerMetadata.getNumDocs(), column); + assertEquals(lazyMetadata.getCardinality(), eagerMetadata.getCardinality(), column); + assertEquals(lazyMetadata.isSorted(), eagerMetadata.isSorted(), column); + assertEquals(lazyMetadata.getMinValue(), eagerMetadata.getMinValue(), column); + assertEquals(lazyMetadata.getMaxValue(), eagerMetadata.getMaxValue(), column); + for (IndexType indexType : indexTypes) { + assertEquals(lazyDataSource.getIndex(indexType) != null, eagerDataSource.getIndex(indexType) != null, + column + " " + indexType); + } + try (PinotSegmentColumnReader eagerReader = new PinotSegmentColumnReader(eager, column); + PinotSegmentColumnReader lazyReader = new PinotSegmentColumnReader(lazy, column)) { + for (int docId = 0; docId < 20; docId++) { + assertTrue(Objects.deepEquals(lazyReader.getValue(docId), eagerReader.getValue(docId)), + column + " docId " + docId); + } + } + } + } finally { + lazy.destroy(); + eager.destroy(); + } + } + private void verifyIndexDirIsV3(File indexDir) { File[] files = indexDir.listFiles(); assertNotNull(files); diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/FieldIndexConfigs.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/FieldIndexConfigs.java index 1d562719a043..6c9c511fab2a 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/FieldIndexConfigs.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/FieldIndexConfigs.java @@ -58,6 +58,24 @@ public Map unwrapIndexes() { .collect(Collectors.toMap(entry -> entry.getKey().getId(), serializer)); } + /// Two instances are equal when they declare the same config for the same index types: index types compare by + /// identity (they are singletons) and index configs by value. + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof FieldIndexConfigs)) { + return false; + } + return _configMap.equals(((FieldIndexConfigs) o)._configMap); + } + + @Override + public int hashCode() { + return _configMap.hashCode(); + } + @Override public String toString() { try { diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java index d31a99d79451..6f90c8ed9280 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java @@ -107,6 +107,11 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig public static final String DISABLE_DIMENSION_TABLE_PRELOAD = "disable.dimension.table.preload"; private static final boolean DEFAULT_DISABLE_DIMENSION_TABLE_PRELOAD = false; + // Whether to create the index container and data source of a physical column on first access instead of for every + // column at segment load. Off by default. See InstanceDataManagerConfig#isLazyColumnMaterialization(). + public static final String LAZY_COLUMN_MATERIALIZATION = "segment.lazy.column.materialization"; + private static final boolean DEFAULT_LAZY_COLUMN_MATERIALIZATION = false; + // To preload segments of table using upsert in parallel for fast upsert metadata recovery. private static final String MAX_SEGMENT_PRELOAD_THREADS = "max.segment.preload.threads"; @@ -357,4 +362,9 @@ public boolean isDimensionTablePreloadDisabled() { return _serverConfig.getProperty(DISABLE_DIMENSION_TABLE_PRELOAD, DEFAULT_DISABLE_DIMENSION_TABLE_PRELOAD); } + + @Override + public boolean isLazyColumnMaterialization() { + return _serverConfig.getProperty(LAZY_COLUMN_MATERIALIZATION, DEFAULT_LAZY_COLUMN_MATERIALIZATION); + } } diff --git a/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfigTest.java b/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfigTest.java index 9c70581fa243..3ee36323a5c0 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfigTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfigTest.java @@ -26,10 +26,22 @@ import static org.apache.pinot.spi.utils.CommonConstants.Server.INSTANCE_ID; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; public class HelixInstanceDataManagerConfigTest { + @Test + public void testLazyColumnMaterializationIsOffByDefault() throws ConfigurationException { + Map props = new HashMap<>(); + props.put(INSTANCE_ID, "testInstance"); + assertFalse(new HelixInstanceDataManagerConfig(new PinotConfiguration(props)).isLazyColumnMaterialization()); + + props.put(HelixInstanceDataManagerConfig.LAZY_COLUMN_MATERIALIZATION, "true"); + assertTrue(new HelixInstanceDataManagerConfig(new PinotConfiguration(props)).isLazyColumnMaterialization()); + } + @Test public void testTierConfigsWithMultipleTiers() throws ConfigurationException { Map props = new HashMap<>(); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java index dfe37da0ea32..90cfef534243 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java @@ -92,4 +92,11 @@ public interface InstanceDataManagerConfig { boolean shouldCheckCRCOnSegmentLoad(); boolean isDimensionTablePreloadDisabled(); + + /// Whether immutable segments create the index container and data source of a physical column on its first access + /// instead of for every column at load. Off by default: it trades per-column heap on wide segments for reporting an + /// unreadable index on first access rather than at load (see the segment loader for the full trade-off). + default boolean isLazyColumnMaterialization() { + return false; + } }