This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/data-3221-9-column-arrays in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 2db3a4c08b22930a381982b5c78c9ad89bb53e16 Author: Xiang Fu <[email protected]> AuthorDate: Sat Sep 5 19:58:00 2026 -0700 DATA-3221 (10): hold segment column metadata in sorted arrays, derive the map on demand A server retains one SegmentMetadataImpl per loaded segment for the segment's lifetime, and it held its columns in a TreeMap. On a 1000-column external-table segment that is a red-black-tree node per column: 40.0 B/column measured, pure bookkeeping on top of the ColumnMetadata the node points at. Hold the columns as two parallel arrays instead -- the names in natural order and their metadata at the same index -- which costs 8.1 B/column, so ~32 B/column less. Lookups binary-search the name array and getAllColumns() is an unmodifiable view of it (SortedStringArraySet, a NavigableSet over a sorted String[] range). getColumnMetadataMap() keeps working and keeps returning a TreeMap, but it is now derived from the arrays on the first call and cached until the columns change, exactly as the per-segment Schema already was. Nothing on the load or query path asks for it any more: SegmentMetadata gains getNumColumns(), getAllColumnMetadata(), forEachColumn(BiConsumer) and addColumnMetadata(String, ColumnMetadata) as additive default methods, and the OSS callers are migrated onto those plus the existing getAllColumns()/getColumnMetadataFor(). LoaderTest asserts the map view is still unbuilt after ImmutableSegmentLoader.load, next to the same assertion for the schema, so a stray caller is caught in CI. Two behaviour notes: - getAllColumns() is now an unmodifiable snapshot rather than the TreeMap's live navigableKeySet. TablesResource#getSegmentMetadata was retaining that live view as its running column intersection and calling retainAll on it, i.e. removing columns from the first segment's own metadata; it now intersects a copy. - getColumnMetadataMap() still returns null for a CONSUMING segment, and writes to the returned map no longer reach the metadata (use addColumnMetadata/removeColumn). Mock-based tests that stubbed only getColumnMetadataMap() now stub the accessors the production code reads. Co-Authored-By: Claude Opus 5 <[email protected]> --- .../metadata/segment/SegmentZKMetadataUtils.java | 10 +- .../controller/utils/SegmentMetadataMockUtils.java | 27 ++- .../dedup/BasePartitionDedupMetadataManager.java | 2 +- .../indexsegment/immutable/EmptyIndexSegment.java | 7 +- .../immutable/ImmutableSegmentImpl.java | 53 ++--- .../immutable/ImmutableSegmentLoader.java | 28 +-- .../immutable/PhysicalColumnNames.java | 37 ++- .../segment/index/loader/IndexLoadingConfig.java | 6 +- .../defaultcolumn/BaseDefaultColumnHandler.java | 2 +- .../LegacyRawValueInvertedIndexCleanup.java | 2 +- .../PartitionIdVirtualColumnProvider.java | 13 +- .../upsert/BasePartitionUpsertMetadataManager.java | 2 +- ...apPartitionDedupMetadataManagerWithTTLTest.java | 5 +- .../local/dedup/DedupWithTimestampColumnTest.java | 7 +- .../immutable/EmptyIndexSegmentTest.java | 2 +- .../immutable/ImmutableSegmentImplTest.java | 5 +- .../immutable/MockSegmentMetadata.java | 55 +++++ .../segment/index/SegmentMetadataImplTest.java | 101 ++++++++ .../local/segment/index/loader/LoaderTest.java | 4 + ...ertMetadataManagerForConsistentDeletesTest.java | 5 +- ...rrentMapPartitionUpsertMetadataManagerTest.java | 10 +- .../apache/pinot/segment/spi/SegmentMetadata.java | 34 +++ .../spi/index/metadata/SegmentMetadataImpl.java | 228 ++++++++++++++---- .../spi/index/metadata/SortedStringArraySet.java | 264 +++++++++++++++++++++ .../index/metadata/SortedStringArraySetTest.java | 106 +++++++++ .../resources/SegmentCompressionStatsReader.java | 2 +- .../pinot/server/api/resources/TablesResource.java | 10 +- .../SegmentCompressionStatsReaderTest.java | 8 +- 28 files changed, 862 insertions(+), 173 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtils.java index 93a2ad0bbc1..a6b291869ab 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtils.java @@ -167,16 +167,14 @@ public class SegmentZKMetadataUtils { // Set partition metadata Map<String, ColumnPartitionMetadata> columnPartitionMap = new HashMap<>(); - for (Map.Entry<String, ColumnMetadata> entry : segmentMetadata.getColumnMetadataMap().entrySet()) { - ColumnMetadata columnMetadata = entry.getValue(); + segmentMetadata.forEachColumn((column, columnMetadata) -> { PartitionFunction partitionFunction = columnMetadata.getPartitionFunction(); if (partitionFunction != null) { - ColumnPartitionMetadata columnPartitionMetadata = + columnPartitionMap.put(column, new ColumnPartitionMetadata(partitionFunction.getName(), partitionFunction.getNumPartitions(), - columnMetadata.getPartitions(), partitionFunction.getFunctionConfig()); - columnPartitionMap.put(entry.getKey(), columnPartitionMetadata); + columnMetadata.getPartitions(), partitionFunction.getFunctionConfig())); } - } + }); segmentZKMetadata.setPartitionMetadata( !columnPartitionMap.isEmpty() ? new SegmentPartitionMetadata(columnPartitionMap) : null); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/utils/SegmentMetadataMockUtils.java b/pinot-controller/src/test/java/org/apache/pinot/controller/utils/SegmentMetadataMockUtils.java index 532b5dc2b0f..ee26025fa1a 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/utils/SegmentMetadataMockUtils.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/utils/SegmentMetadataMockUtils.java @@ -22,6 +22,7 @@ import java.util.Set; import java.util.TreeMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BiConsumer; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.partition.function.MurmurPartitionFunction; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -30,6 +31,9 @@ import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.joda.time.Interval; import org.mockito.Mockito; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -40,6 +44,22 @@ public class SegmentMetadataMockUtils { private SegmentMetadataMockUtils() { } + /// Stubs the column accessors of a mocked segment metadata. Mockito stubs every method, so stubbing only + /// `getColumnMetadataMap()` leaves the accessors production code reads (the real implementation holds sorted + /// arrays, not a map) answering `null`. + private static void stubColumns(SegmentMetadata segmentMetadata, TreeMap<String, ColumnMetadata> columns) { + when(segmentMetadata.getColumnMetadataMap()).thenReturn(columns); + when(segmentMetadata.getAllColumns()).thenReturn(columns.navigableKeySet()); + when(segmentMetadata.getAllColumnMetadata()).thenReturn(columns.values()); + when(segmentMetadata.getNumColumns()).thenReturn(columns.size()); + when(segmentMetadata.getColumnMetadataFor(anyString())).thenAnswer( + call -> columns.get(call.<String>getArgument(0))); + doAnswer(call -> { + columns.forEach(call.<BiConsumer<String, ColumnMetadata>>getArgument(0)); + return null; + }).when(segmentMetadata).forEachColumn(any()); + } + public static SegmentMetadata mockSegmentMetadata(String tableName, String segmentName, int numTotalDocs, String crc, long startTime, long endTime, TimeUnit timeUnit) { SegmentMetadata segmentMetadata = Mockito.mock(SegmentMetadata.class); @@ -100,7 +120,7 @@ public class SegmentMetadataMockUtils { when(colMeta.getPartitionFunction()).thenReturn(new MurmurPartitionFunction(numPartitions, null)); TreeMap<String, ColumnMetadata> columnMetadataMap = new TreeMap<>(); columnMetadataMap.put(partitionColumn, colMeta); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + stubColumns(segmentMetadata, columnMetadataMap); return segmentMetadata; } @@ -112,9 +132,6 @@ public class SegmentMetadataMockUtils { when(columnMetadata.getPartitionFunction()).thenReturn(new MurmurPartitionFunction(5, null)); SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); - if (columnName != null) { - when(segmentMetadata.getColumnMetadataFor(columnName)).thenReturn(columnMetadata); - } when(segmentMetadata.getTableName()).thenReturn(rawTableName); when(segmentMetadata.getName()).thenReturn(segmentName); when(segmentMetadata.getCrc()).thenReturn("0"); @@ -122,7 +139,7 @@ public class SegmentMetadataMockUtils { TreeMap<String, ColumnMetadata> columnMetadataMap = new TreeMap<>(); columnMetadataMap.put(columnName, columnMetadata); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + stubColumns(segmentMetadata, columnMetadataMap); return segmentMetadata; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/dedup/BasePartitionDedupMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/dedup/BasePartitionDedupMetadataManager.java index 693c085d91a..7c0316613a3 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/dedup/BasePartitionDedupMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/dedup/BasePartitionDedupMetadataManager.java @@ -312,7 +312,7 @@ public abstract class BasePartitionDedupMetadataManager implements PartitionDedu // so far to process this segment, as mutable segment is always considered to be within TTL return _largestSeenTime.get(); } - return ((Number) segment.getSegmentMetadata().getColumnMetadataMap().get(_dedupTimeColumn) + return ((Number) segment.getSegmentMetadata().getColumnMetadataFor(_dedupTimeColumn) .getMaxValue()).doubleValue(); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java index 8d99954ff54..672538180e5 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java @@ -19,7 +19,6 @@ package org.apache.pinot.segment.local.indexsegment.immutable; import com.google.common.base.Preconditions; -import java.util.Collections; import java.util.List; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; @@ -80,16 +79,16 @@ public class EmptyIndexSegment implements ImmutableSegment { return _segmentMetadata; } - // Both are views of the column metadata map, so neither builds the segment schema (see SegmentMetadataImpl) + // Both are views of the segment's column metadata, so neither builds the segment schema (see SegmentMetadataImpl) @Override public Set<String> getColumnNames() { - return Collections.unmodifiableSet(_segmentMetadata.getColumnMetadataMap().keySet()); + return _segmentMetadata.getAllColumns(); } @Override public Set<String> getPhysicalColumnNames() { - return new PhysicalColumnNames(_segmentMetadata.getColumnMetadataMap()); + return new PhysicalColumnNames(_segmentMetadata); } @Override 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 7d29a1a9d1d..dd4ef193c99 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 @@ -25,13 +25,11 @@ import java.io.FileOutputStream; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collections; import java.util.HashMap; import java.util.HashSet; 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.atomic.AtomicBoolean; @@ -156,19 +154,15 @@ public class ImmutableSegmentImpl implements ImmutableSegment { _columnMaterializer = null; _openStructChildren = null; _materializationLock = null; - TreeMap<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); - _columnNames = Collections.unmodifiableSet(columnMetadataMap.keySet()); - _physicalColumnNames = new PhysicalColumnNames(columnMetadataMap); - _dataSources = new Object2ObjectOpenHashMap<>(columnMetadataMap.size()); + _columnNames = segmentMetadata.getAllColumns(); + _physicalColumnNames = new PhysicalColumnNames(segmentMetadata); + _dataSources = new Object2ObjectOpenHashMap<>(segmentMetadata.getNumColumns()); Map<String, Map<String, DataSource>> openStructDenseChildren = new HashMap<>(); Map<String, DataSource> openStructSparseChildren = new HashMap<>(); Set<String> openStructParents = new HashSet<>(); - for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.entrySet()) { - String colName = entry.getKey(); - ColumnMetadata columnMetadata = entry.getValue(); - + segmentMetadata.forEachColumn((colName, columnMetadata) -> { if (columnMetadata instanceof ColumnMetadataImpl && ((ColumnMetadataImpl) columnMetadata).isMaterializedChild()) { String parent = ((ColumnMetadataImpl) columnMetadata).getParentColumn(); openStructParents.add(parent); @@ -179,7 +173,7 @@ public class ImmutableSegmentImpl implements ImmutableSegment { openStructDenseChildren.computeIfAbsent(parent, k -> new HashMap<>()) .put(OpenStructNaming.parseKey(colName), childDs); } - continue; + return; } if (columnMetadata.getFieldSpec().getDataType() == FieldSpec.DataType.MAP) { @@ -187,11 +181,11 @@ public class ImmutableSegmentImpl implements ImmutableSegment { } else { _dataSources.put(colName, new ImmutableDataSource(columnMetadata, _indexContainerMap.get(colName))); } - } + }); for (String parent : openStructParents) { // The parent's spec comes from its column metadata, not from the segment schema (see _columnNames) - ColumnMetadata parentMetadata = columnMetadataMap.get(parent); + ColumnMetadata parentMetadata = segmentMetadata.getColumnMetadataFor(parent); FieldSpec fieldSpec = parentMetadata != null ? parentMetadata.getFieldSpec() : null; if (!(fieldSpec instanceof ComplexFieldSpec)) { continue; @@ -231,9 +225,8 @@ public class ImmutableSegmentImpl implements ImmutableSegment { _columnMaterializer = columnMaterializer; _openStructChildren = groupOpenStructChildren(segmentMetadata); _materializationLock = new ReentrantReadWriteLock(); - TreeMap<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); - _columnNames = Collections.unmodifiableSet(columnMetadataMap.keySet()); - _physicalColumnNames = new PhysicalColumnNames(columnMetadataMap); + _columnNames = segmentMetadata.getAllColumns(); + _physicalColumnNames = new PhysicalColumnNames(segmentMetadata); _dataSources = new ConcurrentHashMap<>(); for (String column : materializedIndexContainers.keySet()) { materializeDataSource(column); @@ -244,21 +237,14 @@ public class ImmutableSegmentImpl implements ImmutableSegment { /// metadata declares them complex (the same rule the eager constructor applies). @Nullable private static Map<String, List<String>> groupOpenStructChildren(SegmentMetadataImpl segmentMetadata) { - Map<String, List<String>> children = null; - Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); - for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.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()); + Map<String, List<String>> children = new HashMap<>(); + segmentMetadata.forEachColumn((column, columnMetadata) -> { + if (columnMetadata instanceof ColumnMetadataImpl impl && impl.isMaterializedChild()) { + children.computeIfAbsent(impl.getParentColumn(), k -> new ArrayList<>()).add(column); } - } - if (children == null) { - return null; - } + }); children.keySet().removeIf(parent -> { - ColumnMetadata parentMetadata = columnMetadataMap.get(parent); + ColumnMetadata parentMetadata = segmentMetadata.getColumnMetadataFor(parent); return parentMetadata == null || !(parentMetadata.getFieldSpec() instanceof ComplexFieldSpec); }); return children.isEmpty() ? null : children; @@ -268,7 +254,7 @@ public class ImmutableSegmentImpl implements ImmutableSegment { /// 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); + ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(column); boolean openStructParent = _openStructChildren != null && _openStructChildren.containsKey(column); if (!openStructParent && (columnMetadata == null || isMaterializedChild(columnMetadata))) { return null; @@ -294,11 +280,10 @@ public class ImmutableSegmentImpl implements ImmutableSegment { } 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); + ColumnMetadata childMetadata = _segmentMetadata.getColumnMetadataFor(child); DataSource childDataSource = new ImmutableDataSource(childMetadata, materializedIndexContainer(child, childMetadata)); if (OpenStructNaming.isSparseColumn(child)) { @@ -307,7 +292,7 @@ public class ImmutableSegmentImpl implements ImmutableSegment { denseChildren.put(OpenStructNaming.parseKey(child), childDataSource); } } - ColumnMetadata parentMetadata = columnMetadataMap.get(parent); + ColumnMetadata parentMetadata = _segmentMetadata.getColumnMetadataFor(parent); ComplexFieldSpec fieldSpec = (ComplexFieldSpec) parentMetadata.getFieldSpec(); List<String> sparseKeys = parentMetadata instanceof ColumnMetadataImpl impl ? impl.getSparseKeys() : null; return new ImmutableOpenStructDataSource(fieldSpec, denseChildren, sparseChild, _segmentMetadata.getTotalDocs(), @@ -412,7 +397,7 @@ public class ImmutableSegmentImpl implements ImmutableSegment { 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); + ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(column); if (columnMetadata != null) { Lock lock = _materializationLock.readLock(); lock.lock(); 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 8f980ebd13a..fc08cf71879 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 @@ -39,7 +39,6 @@ import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProvide import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProviderFactory; import org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexContainer; import org.apache.pinot.segment.local.utils.SegmentOperationsThrottlerSet; -import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.ImmutableSegment; import org.apache.pinot.segment.spi.converter.SegmentFormatConverter; import org.apache.pinot.segment.spi.creator.SegmentVersion; @@ -211,9 +210,8 @@ public class ImmutableSegmentLoader { } // Remove columns not in schema from the metadata - Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); if (schema != null) { - Set<String> columnsInMetadata = new HashSet<>(columnMetadataMap.keySet()); + Set<String> columnsInMetadata = new HashSet<>(segmentMetadata.getAllColumns()); columnsInMetadata.removeIf(schema::hasColumn); // Materialized OPEN_STRUCT child columns (col$key, col$__sparse__) live in segment metadata // but not in the user-facing schema. Keep them when the parent OPEN_STRUCT column is in the @@ -233,7 +231,7 @@ public class ImmutableSegmentLoader { } } } else { - indexLoadingConfig.addKnownColumns(columnMetadataMap.keySet()); + indexLoadingConfig.addKnownColumns(segmentMetadata.getAllColumns()); } SegmentDirectory.Reader segmentReader = segmentDirectory.createReader(); @@ -245,11 +243,13 @@ public class ImmutableSegmentLoader { return segment; } - Map<String, ColumnIndexContainer> indexContainerMap = new Object2ObjectOpenHashMap<>(columnMetadataMap.size()); - for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.entrySet()) { + Map<String, ColumnIndexContainer> indexContainerMap = + new Object2ObjectOpenHashMap<>(segmentMetadata.getNumColumns()); + for (String column : segmentMetadata.getAllColumns()) { // FIXME: text-index only works with local SegmentDirectory - indexContainerMap.put(entry.getKey(), - new PhysicalColumnIndexContainer(segmentReader, entry.getValue(), indexLoadingConfig)); + indexContainerMap.put(column, + new PhysicalColumnIndexContainer(segmentReader, segmentMetadata.getColumnMetadataFor(column), + indexLoadingConfig)); } instantiateVirtualColumns(segmentMetadata, indexContainerMap); @@ -287,14 +287,13 @@ public class ImmutableSegmentLoader { SegmentDirectory.Reader segmentReader, SegmentMetadataImpl segmentMetadata, IndexLoadingConfig indexLoadingConfig) throws IOException { - Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); MultiColumnLuceneTextIndexReader mcTextReader = null; Set<String> mcTextColumns = Set.of(); if (segmentReader.hasMultiColumnTextIndex()) { mcTextReader = new MultiColumnLuceneTextIndexReader(segmentMetadata); mcTextColumns = Set.copyOf(segmentMetadata.getMultiColumnTextMetadata().getColumns()); } - ColumnMaterializer columnMaterializer = new ColumnMaterializer(segmentReader, columnMetadataMap.keySet(), + ColumnMaterializer columnMaterializer = new ColumnMaterializer(segmentReader, segmentMetadata.getAllColumns(), indexLoadingConfig.getFieldIndexConfigByColName(), indexLoadingConfig.isForwardIndexOnly(), mcTextReader, mcTextColumns); @@ -305,7 +304,7 @@ public class ImmutableSegmentLoader { if (segmentReader.hasStarTreeIndex()) { starTreeIndexContainer = new StarTreeIndexContainer(segmentReader, segmentMetadata, column -> indexContainerMap.computeIfAbsent(column, - k -> columnMaterializer.createIndexContainer(columnMetadataMap.get(k)))); + k -> columnMaterializer.createIndexContainer(segmentMetadata.getColumnMetadataFor(k)))); } return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata, columnMaterializer, indexContainerMap, @@ -314,16 +313,15 @@ public class ImmutableSegmentLoader { /// Creates the index containers and column metadata of the built-in virtual columns and registers them in the /// segment metadata. Registering the metadata is what makes the segment schema include the virtual columns: the - /// schema is derived from the column metadata map on demand ([SegmentMetadataImpl#getSchema()]) and is deliberately + /// schema is derived from the column metadata on demand ([SegmentMetadataImpl#getSchema()]) and is deliberately /// not built here, so a loaded segment retains no per-column schema entries until something asks for its schema. /// A physical column of the same name wins, as in the schema-based registration this replaces. private static void instantiateVirtualColumns(SegmentMetadataImpl segmentMetadata, Map<String, ColumnIndexContainer> indexContainerMap) { - Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); String segmentName = segmentMetadata.getName(); for (BuiltInVirtualColumnDefinitions.Definition definition : BuiltInVirtualColumnDefinitions.DEFINITIONS) { String columnName = definition.getName(); - if (columnMetadataMap.containsKey(columnName)) { + if (segmentMetadata.getColumnMetadataFor(columnName) != null) { continue; } FieldSpec fieldSpec = VirtualColumnProviderFactory.createBuiltInFieldSpec(definition, segmentName); @@ -331,7 +329,7 @@ public class ImmutableSegmentLoader { new VirtualColumnContext(fieldSpec, segmentMetadata.getTotalDocs(), segmentMetadata); VirtualColumnProvider provider = VirtualColumnProviderFactory.buildProvider(context); indexContainerMap.put(columnName, provider.buildColumnIndexContainer(context)); - columnMetadataMap.put(columnName, provider.buildMetadata(context)); + segmentMetadata.addColumnMetadata(columnName, provider.buildMetadata(context)); } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java index b9dce7f3b86..3ca7967ecbc 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java @@ -20,29 +20,28 @@ package org.apache.pinot.segment.local.indexsegment.immutable; import java.util.AbstractSet; import java.util.Iterator; -import java.util.Map; import java.util.NoSuchElementException; -import java.util.SortedMap; import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.SegmentMetadata; -/// Unmodifiable view of the physical columns of an immutable segment: the keys of its column metadata map whose field -/// spec is not produced by a virtual column provider, in the map's (sorted) key order. +/// Unmodifiable view of the physical columns of an immutable segment: the columns of its metadata whose field spec is +/// not produced by a virtual column provider, in natural column-name order. /// /// It is a view rather than a copy so a segment retains nothing per column for it: the segment schema this replaces -/// held a `TreeMap` entry per column, and a cached `TreeSet` would hold the same. `contains` is one map lookup and -/// iteration is a filtered pass over the map. The virtual column count is taken once at construction, which is sound -/// because the column metadata map is fixed once the segment is loaded. +/// held a `TreeMap` entry per column, and a cached `TreeSet` would hold the same. `contains` is one column lookup and +/// iteration is a filtered pass over the column metadata. The virtual column count is taken once at construction, +/// which is sound because the column metadata is fixed once the segment is loaded. /// -/// Thread-safe for reads, like the underlying map once loaded. +/// Thread-safe for reads, like the underlying segment metadata once loaded. final class PhysicalColumnNames extends AbstractSet<String> { - private final SortedMap<String, ColumnMetadata> _columnMetadataMap; + private final SegmentMetadata _segmentMetadata; private final int _numVirtualColumns; - PhysicalColumnNames(SortedMap<String, ColumnMetadata> columnMetadataMap) { - _columnMetadataMap = columnMetadataMap; + PhysicalColumnNames(SegmentMetadata segmentMetadata) { + _segmentMetadata = segmentMetadata; int numVirtualColumns = 0; - for (ColumnMetadata columnMetadata : columnMetadataMap.values()) { + for (ColumnMetadata columnMetadata : segmentMetadata.getAllColumnMetadata()) { if (!isPhysical(columnMetadata)) { numVirtualColumns++; } @@ -59,26 +58,26 @@ final class PhysicalColumnNames extends AbstractSet<String> { if (!(o instanceof String)) { return false; } - ColumnMetadata columnMetadata = _columnMetadataMap.get(o); + ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor((String) o); return columnMetadata != null && isPhysical(columnMetadata); } @Override public int size() { - return _columnMetadataMap.size() - _numVirtualColumns; + return _segmentMetadata.getNumColumns() - _numVirtualColumns; } @Override public Iterator<String> iterator() { - Iterator<Map.Entry<String, ColumnMetadata>> entries = _columnMetadataMap.entrySet().iterator(); + Iterator<ColumnMetadata> columnMetadata = _segmentMetadata.getAllColumnMetadata().iterator(); return new Iterator<>() { private String _next = advance(); private String advance() { - while (entries.hasNext()) { - Map.Entry<String, ColumnMetadata> entry = entries.next(); - if (isPhysical(entry.getValue())) { - return entry.getKey(); + while (columnMetadata.hasNext()) { + ColumnMetadata next = columnMetadata.next(); + if (isPhysical(next)) { + return next.getColumnName(); } } return null; 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 c7bd326e65d..416f1418aaa 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 @@ -29,7 +29,6 @@ import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.segment.local.segment.index.loader.columnminmaxvalue.ColumnMinMaxValueGeneratorMode; import org.apache.pinot.segment.local.utils.TableConfigUtils; -import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; @@ -427,8 +426,7 @@ public class IndexLoadingConfig { if (_indexConfigsByColName == null || _dirty) { refreshIndexConfigs(); } - for (Map.Entry<String, ColumnMetadata> entry : segmentMetadata.getColumnMetadataMap().entrySet()) { - String childColumn = entry.getKey(); + for (String childColumn : segmentMetadata.getAllColumns()) { if (!childColumn.contains(OpenStructNaming.SEPARATOR) || _indexConfigsByColName.containsKey(childColumn)) { continue; } @@ -450,7 +448,7 @@ public class IndexLoadingConfig { if (keyFieldConfig == null) { keyFieldConfig = openStructConfig.getDefaultValueFieldConfig(); } - FieldSpec childFieldSpec = entry.getValue().getFieldSpec(); + FieldSpec childFieldSpec = segmentMetadata.getColumnMetadataFor(childColumn).getFieldSpec(); boolean enableInverted = openStructConfig.shouldEnableInvertedIndexForKey(key); FieldIndexConfigs childConfigs = new FieldIndexConfigs.Builder( FieldIndexConfigsUtil.fromFieldConfig(keyFieldConfig, childFieldSpec)) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java index cd29c2045b0..96997664cf1 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java @@ -336,7 +336,7 @@ public abstract class BaseDefaultColumnHandler implements DefaultColumnHandler { } // Compute REMOVE actions. - for (ColumnMetadata columnMetadata : _segmentMetadata.getColumnMetadataMap().values()) { + for (ColumnMetadata columnMetadata : _segmentMetadata.getAllColumnMetadata()) { String column = columnMetadata.getColumnName(); // Only remove auto-generated columns if (!_schema.hasColumn(column) && columnMetadata.isAutoGenerated()) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/LegacyRawValueInvertedIndexCleanup.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/LegacyRawValueInvertedIndexCleanup.java index c81c9a7de86..0ed08889ba6 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/LegacyRawValueInvertedIndexCleanup.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/LegacyRawValueInvertedIndexCleanup.java @@ -78,7 +78,7 @@ public final class LegacyRawValueInvertedIndexCleanup { throws IOException { SegmentMetadataImpl segmentMetadata = (SegmentMetadataImpl) segmentWriter.toSegmentDirectory().getSegmentMetadata(); String segmentName = segmentMetadata.getName(); - for (ColumnMetadata columnMetadata : segmentMetadata.getColumnMetadataMap().values()) { + for (ColumnMetadata columnMetadata : segmentMetadata.getAllColumnMetadata()) { String column = columnMetadata.getColumnName(); if (!segmentWriter.hasIndexFor(column, StandardIndexes.inverted())) { continue; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/virtualcolumn/PartitionIdVirtualColumnProvider.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/virtualcolumn/PartitionIdVirtualColumnProvider.java index 997ff24ff13..49931993e43 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/virtualcolumn/PartitionIdVirtualColumnProvider.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/virtualcolumn/PartitionIdVirtualColumnProvider.java @@ -24,11 +24,9 @@ import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; -import java.util.Map; import java.util.Set; import org.apache.pinot.segment.local.segment.index.readers.BaseImmutableDictionary; import org.apache.pinot.segment.local.segment.index.readers.constant.ConstantMVInvertedIndexReader; -import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.SegmentMetadata; import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl; import org.apache.pinot.segment.spi.index.reader.Dictionary; @@ -88,20 +86,17 @@ public class PartitionIdVirtualColumnProvider implements VirtualColumnProvider { List<String> partitionInfo = new ArrayList<>(); SegmentMetadata segmentMetadata = context.getSegmentMetadata(); - if (segmentMetadata != null && segmentMetadata.getColumnMetadataMap() != null) { + if (segmentMetadata != null) { // Get partition info from all partitioned columns in the segment metadata - Map<String, ColumnMetadata> columnMetadataMap = segmentMetadata.getColumnMetadataMap(); - for (Map.Entry<String, ColumnMetadata> entry : columnMetadataMap.entrySet()) { - String columnName = entry.getKey(); - ColumnMetadata columnMetadata = entry.getValue(); + segmentMetadata.forEachColumn((columnName, columnMetadata) -> { Set<Integer> partitions = columnMetadata.getPartitions(); - if (partitions != null && !partitions.isEmpty()) { + if (partitions != null) { // Add all partition IDs for this column for (Integer partitionId : partitions) { partitionInfo.add(columnName + "_" + partitionId); } } - } + }); } // Ensure we always have at least one entry for multi-value columns diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java index 166ff082500..473a1b604c8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java @@ -317,7 +317,7 @@ public abstract class BasePartitionUpsertMetadataManager implements PartitionUps } protected double getMaxComparisonValue(IndexSegment segment) { - return ((Number) segment.getSegmentMetadata().getColumnMetadataMap().get(_comparisonColumns.get(0)) + return ((Number) segment.getSegmentMetadata().getColumnMetadataFor(_comparisonColumns.get(0)) .getMaxValue()).doubleValue(); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/ConcurrentMapPartitionDedupMetadataManagerWithTTLTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/ConcurrentMapPartitionDedupMetadataManagerWithTTLTest.java index 678d0064af6..c287e0f558c 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/ConcurrentMapPartitionDedupMetadataManagerWithTTLTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/ConcurrentMapPartitionDedupMetadataManagerWithTTLTest.java @@ -22,7 +22,6 @@ import java.io.File; import java.io.IOException; import java.util.Iterator; import java.util.List; -import java.util.TreeMap; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.segment.local.data.manager.TableDataManager; @@ -446,9 +445,7 @@ public class ConcurrentMapPartitionDedupMetadataManagerWithTTLTest { IndexSegment segment = DedupTestUtils.mockSegment(1, 10); SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); ColumnMetadata columnMetadata = mock(ColumnMetadata.class); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>() {{ - this.put(DEDUP_TIME_COLUMN_NAME, columnMetadata); - }}); + when(segmentMetadata.getColumnMetadataFor(DEDUP_TIME_COLUMN_NAME)).thenReturn(columnMetadata); doReturn(System.currentTimeMillis()).when(columnMetadata).getMaxValue(); when(segment.getSegmentMetadata()).thenReturn(segmentMetadata); // throws when not stopped diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/DedupWithTimestampColumnTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/DedupWithTimestampColumnTest.java index b49cbdee78f..e8ec98777be 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/DedupWithTimestampColumnTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/dedup/DedupWithTimestampColumnTest.java @@ -22,7 +22,6 @@ import java.io.File; import java.io.IOException; import java.util.Iterator; import java.util.List; -import java.util.TreeMap; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.segment.local.data.manager.TableDataManager; @@ -201,9 +200,7 @@ public class DedupWithTimestampColumnTest { // TIMESTAMP values are stored as LONG (epoch milliseconds) long currentTimeMillis = System.currentTimeMillis(); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>() {{ - this.put(DEDUP_TIME_COLUMN_NAME, columnMetadata); - }}); + when(segmentMetadata.getColumnMetadataFor(DEDUP_TIME_COLUMN_NAME)).thenReturn(columnMetadata); doReturn(currentTimeMillis).when(columnMetadata).getMaxValue(); when(segment.getSegmentMetadata()).thenReturn(segmentMetadata); @@ -214,7 +211,7 @@ public class DedupWithTimestampColumnTest { // Verify the segment was added successfully assertNotNull(segment.getSegmentMetadata()); - assertEquals(segment.getSegmentMetadata().getColumnMetadataMap().get(DEDUP_TIME_COLUMN_NAME).getMaxValue(), + assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(DEDUP_TIME_COLUMN_NAME).getMaxValue(), currentTimeMillis); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java index f308cd6a6ca..1acad42c148 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java @@ -103,7 +103,7 @@ public class EmptyIndexSegmentTest { columnMetadataMap.put(column, new EmptyColumnMetadata(new DimensionFieldSpec(column, FieldSpec.DataType.INT, true), null, null)); } - when(metadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + MockSegmentMetadata.withColumns(metadata, columnMetadataMap); EmptyIndexSegment segment = new EmptyIndexSegment(metadata); assertEquals(new ArrayList<>(segment.getColumnNames()), List.of("a", "b")); 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 41a48241836..21cf674fc29 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 @@ -122,8 +122,7 @@ public class ImmutableSegmentImplTest { private static ImmutableSegmentImpl createSegment(SegmentDirectory segmentDirectory) { SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); when(segmentMetadata.getName()).thenReturn("seg"); - // getColumnMetadataMap() is declared as a TreeMap, so an immutable Map.of() will not do here. - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>()); + MockSegmentMetadata.withColumns(segmentMetadata, Map.of()); return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata, Map.of(), null); } @@ -443,7 +442,7 @@ public class ImmutableSegmentImplTest { for (ColumnMetadata column : columns) { columnMetadataMap.put(column.getColumnName(), column); } - when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + MockSegmentMetadata.withColumns(segmentMetadata, columnMetadataMap); return segmentMetadata; } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/MockSegmentMetadata.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/MockSegmentMetadata.java new file mode 100644 index 00000000000..9d393516eed --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/MockSegmentMetadata.java @@ -0,0 +1,55 @@ +/** + * 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.util.Collections; +import java.util.Map; +import java.util.TreeMap; +import java.util.function.BiConsumer; +import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.when; + + +/// Stubs the column accessors of a mocked [SegmentMetadataImpl]. +/// +/// Mockito stubs every method of the mock, so stubbing [SegmentMetadataImpl#getColumnMetadataMap()] alone leaves the +/// accessors the segment actually reads (the real implementation holds sorted arrays, not a map) answering `null`. +final class MockSegmentMetadata { + private MockSegmentMetadata() { + } + + static SegmentMetadataImpl withColumns(SegmentMetadataImpl metadata, Map<String, ColumnMetadata> columns) { + TreeMap<String, ColumnMetadata> sorted = new TreeMap<>(columns); + when(metadata.getColumnMetadataMap()).thenReturn(sorted); + when(metadata.getAllColumns()).thenReturn(Collections.unmodifiableNavigableSet(sorted.navigableKeySet())); + when(metadata.getAllColumnMetadata()).thenReturn(Collections.unmodifiableCollection(sorted.values())); + when(metadata.getNumColumns()).thenReturn(sorted.size()); + when(metadata.getColumnMetadataFor(anyString())).thenAnswer(call -> sorted.get(call.<String>getArgument(0))); + doAnswer(call -> { + sorted.forEach(call.<BiConsumer<String, ColumnMetadata>>getArgument(0)); + return null; + }).when(metadata).forEachColumn(any()); + return metadata; + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java index c829c36c32a..d6f922fdd23 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java @@ -26,9 +26,12 @@ import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.NavigableSet; import java.util.Set; +import java.util.TreeMap; import java.util.concurrent.TimeUnit; import org.apache.commons.configuration2.ex.ConfigurationException; import org.apache.commons.io.FileUtils; @@ -45,6 +48,7 @@ import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.metadata.ColumnMetadataImpl; +import org.apache.pinot.segment.spi.index.metadata.EmptyColumnMetadata; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.spi.config.table.FieldConfig; @@ -74,6 +78,7 @@ import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNotSame; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; @@ -371,6 +376,102 @@ public class SegmentMetadataImplTest { } } + /// The columns are held as sorted arrays; the `TreeMap` view exists only for compatibility and costs a map entry + /// per column, so it is derived on the first getColumnMetadataMap() and never by the accessors the load and query + /// paths use. + @Test + public void testColumnMetadataMapDerivedLazily() + throws Exception { + long materializations = SegmentMetadataImpl.getNumColumnMetadataMapMaterializations(); + SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory); + assertFalse(metadata.isColumnMetadataMapMaterialized()); + + List<String> columns = new ArrayList<>(metadata.getAllColumns()); + assertEquals(metadata.getNumColumns(), columns.size()); + assertEquals(metadata.getAllColumnMetadata().size(), columns.size()); + for (String column : columns) { + assertNotNull(metadata.getColumnMetadataFor(column), column); + } + Map<String, ColumnMetadata> visited = new LinkedHashMap<>(); + metadata.forEachColumn(visited::put); + assertEquals(new ArrayList<>(visited.keySet()), columns); + metadata.toJson(null); + assertFalse(metadata.isColumnMetadataMapMaterialized(), "reading the columns must not build the map"); + assertEquals(SegmentMetadataImpl.getNumColumnMetadataMapMaterializations(), materializations); + + TreeMap<String, ColumnMetadata> map = metadata.getColumnMetadataMap(); + assertTrue(metadata.isColumnMetadataMapMaterialized()); + assertEquals(SegmentMetadataImpl.getNumColumnMetadataMapMaterializations(), materializations + 1); + assertEquals(map, visited); + assertEquals(new ArrayList<>(map.keySet()), columns); + assertSame(metadata.getColumnMetadataMap(), map); + assertEquals(SegmentMetadataImpl.getNumColumnMetadataMapMaterializations(), materializations + 1); + assertNull(metadata.getColumnMetadataFor("noSuchColumn")); + } + + /// getAllColumns() is a view of the metadata's own name array, so it must refuse every mutator rather than let a + /// caller narrow a loaded segment's columns, and it must not reflect later column changes. + @Test + public void testGetAllColumnsIsAnUnmodifiableSnapshot() + throws Exception { + SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory); + NavigableSet<String> columns = metadata.getAllColumns(); + String column = columns.stream().filter(c -> !c.equals(metadata.getTimeColumn())).findFirst().orElseThrow(); + assertThrows(UnsupportedOperationException.class, () -> columns.remove(column)); + assertThrows(UnsupportedOperationException.class, () -> columns.retainAll(Set.of(column))); + + metadata.removeColumn(column); + assertTrue(columns.contains(column), "the earlier view stays the snapshot it was"); + assertFalse(metadata.getAllColumns().contains(column)); + } + + /// The loader registers the built-in virtual columns through addColumnMetadata(), which has to keep the arrays + /// sorted and drop both derived views. + @Test + public void testAddColumnMetadataKeepsColumnsSortedAndDropsDerivedViews() + throws Exception { + SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory); + assertNotNull(metadata.getColumnMetadataMap()); + assertNotNull(metadata.getSchema()); + int numColumns = metadata.getNumColumns(); + + String column = "$aVirtualColumn"; + ColumnMetadata added = + new EmptyColumnMetadata(new DimensionFieldSpec(column, FieldSpec.DataType.INT, true), null, null); + metadata.addColumnMetadata(column, added); + assertFalse(metadata.isColumnMetadataMapMaterialized()); + assertFalse(metadata.isSchemaMaterialized()); + assertEquals(metadata.getNumColumns(), numColumns + 1); + assertSame(metadata.getColumnMetadataFor(column), added); + assertEquals(new ArrayList<>(metadata.getAllColumns()), new ArrayList<>(metadata.getColumnMetadataMap().keySet())); + assertEquals(metadata.getAllColumns().first(), column, "must be inserted in natural order, not appended"); + assertTrue(metadata.getSchema().hasColumn(column)); + + // Re-registering replaces in place rather than duplicating the column + ColumnMetadata replacement = + new EmptyColumnMetadata(new DimensionFieldSpec(column, FieldSpec.DataType.LONG, true), null, null); + metadata.addColumnMetadata(column, replacement); + assertEquals(metadata.getNumColumns(), numColumns + 1); + assertSame(metadata.getColumnMetadataFor(column), replacement); + } + + /// The metadata JSON is a public REST payload: it must list the columns in the same natural order the map view + /// does, filter included. + @Test + public void testToJsonColumnOrderMatchesTheMapView() + throws Exception { + SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory); + List<String> expected = new ArrayList<>(metadata.getColumnMetadataMap().keySet()); + List<String> actual = new ArrayList<>(); + metadata.toJson(null).get("columns").forEach(column -> actual.add(column.get("columnName").asText())); + assertEquals(actual, expected); + + Set<String> filter = Set.of(expected.get(expected.size() - 1), expected.get(0)); + List<String> filtered = new ArrayList<>(); + metadata.toJson(filter).get("columns").forEach(column -> filtered.add(column.get("columnName").asText())); + assertEquals(filtered, List.of(expected.get(0), expected.get(expected.size() - 1))); + } + /// removeColumn() drops the column from the column metadata and from any schema derived afterwards. @Test public void testRemoveColumnInvalidatesDerivedSchema() 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 a1bf466bdbc..5b4d87068d9 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 @@ -237,6 +237,10 @@ public class LoaderTest { // path would re-inflate the per-column schema footprint of every segment a server loads. assertFalse(((SegmentMetadataImpl) indexSegment.getSegmentMetadata()).isSchemaMaterialized(), "the load path must not build the segment schema"); + // Likewise for the column metadata map: the metadata holds sorted arrays, and the map is derived on demand, so a + // stray getColumnMetadataMap() on the load path would cost a map entry per column of every segment loaded. + assertFalse(((SegmentMetadataImpl) indexSegment.getSegmentMetadata()).isColumnMetadataMapMaterialized(), + "the load path must not build the column metadata map"); // Segment metadata that this segment carries is exposed as a real value, and is not marked null SegmentMetadata segmentMetadata = indexSegment.getSegmentMetadata(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java index 52a688907f8..84c4102ef70 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java @@ -25,7 +25,6 @@ import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.TreeMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; @@ -1245,9 +1244,7 @@ public class ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); ColumnMetadata columnMetadata = mock(ColumnMetadata.class); when(segmentMetadata.getTotalDocs()).thenReturn(deleteFlags.length); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap() {{ - this.put(COMPARISON_COLUMNS.get(0), columnMetadata); - }}); + when(segmentMetadata.getColumnMetadataFor(COMPARISON_COLUMNS.get(0))).thenReturn(columnMetadata); ImmutableSegmentImpl segment = mockImmutableSegmentWithSegmentMetadata(1, new ThreadSafeMutableRoaringBitmap(), null, null, segmentMetadata, diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java index a3dcfb12821..d54dae1e659 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java @@ -972,6 +972,8 @@ public class ConcurrentMapPartitionUpsertMetadataManagerTest { columnMetadataMap.put(PRIMARY_KEY_COLUMNS.get(0), primaryKeyColumnMetadata); columnMetadataMap.put(COMPARISON_COLUMNS.get(0), comparisonColumnMetadata); when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + when(segmentMetadata.getColumnMetadataFor(anyString())).thenAnswer( + call -> columnMetadataMap.get(call.<String>getArgument(0))); when(segment.getSegmentMetadata()).thenReturn(segmentMetadata); return segment; @@ -1060,9 +1062,7 @@ public class ConcurrentMapPartitionUpsertMetadataManagerTest { SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); when(segment.getSegmentMetadata()).thenReturn(segmentMetadata); ColumnMetadata columnMetadata = mock(ColumnMetadata.class); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap() {{ - this.put(comparisonColumns.get(0), columnMetadata); - }}); + when(segmentMetadata.getColumnMetadataFor(comparisonColumns.get(0))).thenReturn(columnMetadata); doReturn(endTime).when(columnMetadata).getMaxValue(); if (snapshot != null) { when(segment.loadDocIdsFromSnapshot(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME)).thenReturn(snapshot); @@ -1903,9 +1903,7 @@ public class ConcurrentMapPartitionUpsertMetadataManagerTest { SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); ColumnMetadata columnMetadata = mock(ColumnMetadata.class); when(segmentMetadata.getTotalDocs()).thenReturn(deleteFlags.length); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap() {{ - this.put(COMPARISON_COLUMNS.get(0), columnMetadata); - }}); + when(segmentMetadata.getColumnMetadataFor(COMPARISON_COLUMNS.get(0))).thenReturn(columnMetadata); ImmutableSegmentImpl segment = mockImmutableSegmentWithSegmentMetadata(1, new ThreadSafeMutableRoaringBitmap(), null, null, segmentMetadata, diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java index c1c3ceb4299..c01f2cf39e6 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java @@ -20,12 +20,14 @@ package org.apache.pinot.segment.spi; import com.fasterxml.jackson.databind.JsonNode; import java.io.File; +import java.util.Collection; import java.util.List; import java.util.Map; import java.util.NavigableSet; import java.util.Set; import java.util.TreeMap; import java.util.concurrent.TimeUnit; +import java.util.function.BiConsumer; import javax.annotation.Nullable; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.multicolumntext.MultiColumnTextMetadata; @@ -120,12 +122,44 @@ public interface SegmentMetadata { return getSchema().getColumnNames(); } + /// Number of columns the segment metadata holds. + default int getNumColumns() { + return getColumnMetadataMap().size(); + } + + /// The column metadata of every column, in the natural column-name order of [#getAllColumns()]. + default Collection<ColumnMetadata> getAllColumnMetadata() { + return getColumnMetadataMap().values(); + } + + /// Applies `action` to every (column name, column metadata) pair, in the natural column-name order of + /// [#getAllColumns()]. + default void forEachColumn(BiConsumer<String, ColumnMetadata> action) { + getColumnMetadataMap().forEach(action); + } + + /// Returns the whole column metadata as a map, for callers that need one. + /// + /// An implementation may hold its columns in a form that costs less than a map entry per column and build this map + /// on demand, so load- and query-path code must not call this: it re-inflates a map entry per column for every + /// segment it touches, and a server keeps that for the segment's lifetime. Read column names through + /// [#getAllColumns()], one column through [#getColumnMetadataFor(String)], all of them through + /// [#getAllColumnMetadata()] or [#forEachColumn(BiConsumer)], and mutate through + /// [#addColumnMetadata(String, ColumnMetadata)] / [#removeColumn(String)] rather than through the returned map, + /// whose writes an implementation is free not to see. TreeMap<String, ColumnMetadata> getColumnMetadataMap(); + /// Returns the metadata of the given column, or `null` if the segment has no such column. + @Nullable default ColumnMetadata getColumnMetadataFor(String column) { return getColumnMetadataMap().get(column); } + /// Registers the metadata of a column, replacing any metadata already registered under the same name. + default void addColumnMetadata(String column, ColumnMetadata columnMetadata) { + getColumnMetadataMap().put(column, columnMetadata); + } + /// Removes a column from the segment metadata. void removeColumn(String column); diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java index 5bdc0057410..9cdb5b17775 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java @@ -31,6 +31,9 @@ import java.io.InputStream; import java.text.DateFormat; import java.text.SimpleDateFormat; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; import java.util.Date; import java.util.HashMap; import java.util.HashSet; @@ -43,6 +46,7 @@ import java.util.TimeZone; import java.util.TreeMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BiConsumer; import javax.annotation.Nullable; import org.apache.commons.configuration2.Configuration; import org.apache.commons.configuration2.PropertiesConfiguration; @@ -76,18 +80,21 @@ import org.slf4j.LoggerFactory; /// Segment metadata parsed from `metadata.properties` (plus `creation.meta` and the v3 `index_map`), or built for a /// CONSUMING segment from an explicit [Schema]. /// -/// The segment [Schema] of a metadata-backed instance is derived from [#getColumnMetadataMap()] on the first -/// [#getSchema()] call and cached; it is not built at load. A server retains one instance per loaded segment for the -/// segment's lifetime, and a Schema costs a `TreeMap` entry plus list slots per column on top of the column metadata -/// that already holds every [org.apache.pinot.spi.data.FieldSpec], so building it eagerly doubled the per-column -/// metadata footprint of a wide segment that is never asked for its schema. Everything on the load and query paths -/// reads the column metadata map (or [#getAllColumns()], a view of its keys) instead. Once the loader has registered -/// the built-in virtual columns in the map, the derived schema includes them, exactly as the eagerly built one did. -/// [#removeColumn(String)] drops the cached schema so it is rebuilt without the column. The explicit-schema -/// constructor keeps the caller's Schema as is. +/// A server retains one instance per loaded segment for the segment's lifetime, so the columns are held as two +/// parallel arrays — the names in natural order and their metadata at the same index — rather than as a map: two +/// array slots per column instead of a red-black-tree node, which on a segment of a thousand columns is the +/// difference between a few kilobytes and tens of kilobytes of pure bookkeeping. Lookups +/// ([#getColumnMetadataFor(String)]) binary-search the name array, and [#getAllColumns()] is a view of it. The +/// [#getColumnMetadataMap()] map and the segment [Schema] are both derived from the arrays only when something asks +/// for them, and cached until the columns change; nothing on the load or query path asks. /// -/// Thread-safe for the schema cache (double-checked on a volatile, so one instance per metadata); the rest is -/// populated at load before the metadata is published. +/// Once the loader has registered the built-in virtual columns through [#addColumnMetadata(String, ColumnMetadata)], +/// the derived schema includes them, exactly as the eagerly built one did. [#removeColumn(String)] and +/// [#addColumnMetadata(String, ColumnMetadata)] replace the arrays and drop both derived views. The explicit-schema +/// constructor keeps the caller's Schema as is and holds no column metadata at all. +/// +/// Thread-safe for the derived schema and map (double-checked on volatiles, so one instance per metadata); the +/// columns themselves are populated at load before the metadata is published. public class SegmentMetadataImpl implements SegmentMetadata { private static final Logger LOGGER = LoggerFactory.getLogger(SegmentMetadataImpl.class); @@ -95,12 +102,27 @@ public class SegmentMetadataImpl implements SegmentMetadata { /// segment's schema unbuilt. private static final AtomicLong NUM_SCHEMA_MATERIALIZATIONS = new AtomicLong(); + /// Number of derived column metadata maps built so far, JVM-wide, for the same reason. + private static final AtomicLong NUM_COLUMN_METADATA_MAP_MATERIALIZATIONS = new AtomicLong(); + private final File _indexDir; - private final TreeMap<String, ColumnMetadata> _columnMetadataMap; + /// Column names in natural order, and their metadata at the same index. Both `null` for a CONSUMING segment, which + /// is constructed with an explicit schema and holds no column metadata. Replaced (never written in place) by + /// [#addColumnMetadata(String, ColumnMetadata)] and [#removeColumn(String)], so a view handed out earlier stays a + /// consistent snapshot, and volatile so a metadata published without other synchronization is seen with its + /// columns. + @Nullable + private volatile String[] _columnNames; + @Nullable + private volatile ColumnMetadata[] _columnMetadata; /// The explicit schema of a CONSUMING segment, or the lazily derived schema of a metadata-backed segment (null - /// until [#getSchema()] builds it, and again after [#removeColumn(String)]). + /// until [#getSchema()] builds it, and again whenever the columns change). @Nullable private volatile Schema _schema; + /// The lazily derived map view of the two column arrays (null until [#getColumnMetadataMap()] builds it, and again + /// whenever the columns change). + @Nullable + private volatile TreeMap<String, ColumnMetadata> _columnMetadataMapView; private String _segmentName; private int _totalDocs; private SegmentVersion _segmentVersion; @@ -133,7 +155,6 @@ public class SegmentMetadataImpl implements SegmentMetadata { public SegmentMetadataImpl(InputStream metadataPropertiesInputStream, InputStream creationMetaInputStream) throws IOException, ConfigurationException { _indexDir = null; - _columnMetadataMap = new TreeMap<>(); PropertiesConfiguration segmentMetadataPropertiesConfiguration = CommonsConfigurationUtils.fromInputStream(metadataPropertiesInputStream); @@ -151,7 +172,6 @@ public class SegmentMetadataImpl implements SegmentMetadata { public SegmentMetadataImpl(File indexDir) throws IOException, ConfigurationException { _indexDir = indexDir; - _columnMetadataMap = new TreeMap<>(); PropertiesConfiguration segmentMetadataPropertiesConfiguration = SegmentMetadataUtils.getPropertiesConfiguration(indexDir); @@ -167,7 +187,6 @@ public class SegmentMetadataImpl implements SegmentMetadata { /// For REALTIME consuming segments. public SegmentMetadataImpl(String rawTableName, String segmentName, Schema schema, long creationTime) { _indexDir = null; - _columnMetadataMap = null; _rawTableName = rawTableName; _segmentName = segmentName; _schema = schema; @@ -254,13 +273,17 @@ public class SegmentMetadataImpl implements SegmentMetadata { addPhysicalColumns(segmentMetadata.getList(Segment.DATETIME_COLUMNS), physicalColumns); addPhysicalColumns(segmentMetadata.getList(Segment.COMPLEX_COLUMNS), physicalColumns); - // Build the column metadata map (the schema is derived from it on demand, see getSchema()). Empty segments use a - // stripped-down [EmptyColumnMetadata] since the shape stats (cardinality, element lengths, etc.) are meaningless - // when there are no rows. + // Build the sorted column arrays (the map view and the schema are derived from them on demand, see + // getColumnMetadataMap() and getSchema()). Empty segments use a stripped-down [EmptyColumnMetadata] since the + // shape stats (cardinality, element lengths, etc.) are meaningless when there are no rows. + String[] columns = physicalColumns.toArray(new String[0]); + Arrays.sort(columns); + ColumnMetadata[] columnMetadata = new ColumnMetadata[columns.length]; + _columnNames = columns; + _columnMetadata = columnMetadata; if (_totalDocs > 0) { - for (String column : physicalColumns) { - _columnMetadataMap.put(column, - ColumnMetadataImpl.fromPropertiesConfiguration(segmentMetadata, _totalDocs, column)); + for (int i = 0; i < columns.length; i++) { + columnMetadata[i] = ColumnMetadataImpl.fromPropertiesConfiguration(segmentMetadata, _totalDocs, columns[i]); } // Load index metadata @@ -278,7 +301,7 @@ public class SegmentMetadataImpl implements SegmentMetadata { String[] parsedKeys = ColumnIndexUtils.parseIndexMapKeys(key, _indexDir.getPath()); if (parsedKeys[2].equals(ColumnIndexUtils.MAP_KEY_NAME_SIZE)) { short indexType = indexService.getNumericId(parsedKeys[1]); - ((ColumnMetadataImpl) _columnMetadataMap.get(parsedKeys[0])).addIndexSize(indexType, + ((ColumnMetadataImpl) getColumnMetadataFor(parsedKeys[0])).addIndexSize(indexType, mapConfig.getLong(key)); } } catch (Exception e) { @@ -288,8 +311,8 @@ public class SegmentMetadataImpl implements SegmentMetadata { } } } else { - for (String column : physicalColumns) { - _columnMetadataMap.put(column, EmptyColumnMetadata.fromPropertiesConfiguration(segmentMetadata, column)); + for (int i = 0; i < columns.length; i++) { + columnMetadata[i] = EmptyColumnMetadata.fromPropertiesConfiguration(segmentMetadata, columns[i]); } } @@ -408,9 +431,9 @@ public class SegmentMetadataImpl implements SegmentMetadata { /// {@inheritDoc} /// - /// For a metadata-backed segment the schema is built from the column metadata map on the first call (one - /// `FieldSpec` per column, the built-in virtual columns included once the loader has registered them) and cached - /// until [#removeColumn(String)]. Nothing on the load or query path should call this: a caller there re-inflates + /// For a metadata-backed segment the schema is built from the column metadata on the first call (one `FieldSpec` + /// per column, the built-in virtual columns included once the loader has registered them) and cached until the + /// columns change. Nothing on the load or query path should call this: a caller there re-inflates /// the per-column schema footprint for every segment it touches. Column names are available through /// [#getAllColumns()] and field specs through [#getColumnMetadataFor(String)]. @Override @@ -431,7 +454,7 @@ public class SegmentMetadataImpl implements SegmentMetadata { private Schema buildSchema() { NUM_SCHEMA_MATERIALIZATIONS.incrementAndGet(); Schema schema = new Schema(); - for (ColumnMetadata columnMetadata : _columnMetadataMap.values()) { + for (ColumnMetadata columnMetadata : _columnMetadata) { schema.addField(columnMetadata.getFieldSpec()); } return schema; @@ -451,11 +474,53 @@ public class SegmentMetadataImpl implements SegmentMetadata { return NUM_SCHEMA_MATERIALIZATIONS.get(); } - /// The keys of the column metadata map, i.e. the same names as `getSchema().getColumnNames()` without building the - /// schema. Falls back to the explicit schema of a CONSUMING segment, which has no column metadata map. + /// An unmodifiable view of the sorted column name array, i.e. the same names as `getSchema().getColumnNames()` + /// without building the schema. Falls back to the explicit schema of a CONSUMING segment, which has no column + /// metadata. The view is a snapshot: it does not reflect columns added or removed after this call. @Override public NavigableSet<String> getAllColumns() { - return _columnMetadataMap != null ? _columnMetadataMap.navigableKeySet() : getSchema().getColumnNames(); + String[] columnNames = _columnNames; + return columnNames != null ? new SortedStringArraySet(columnNames) : getSchema().getColumnNames(); + } + + @Override + public int getNumColumns() { + String[] columnNames = _columnNames; + return columnNames != null ? columnNames.length : getSchema().size(); + } + + /// An unmodifiable view of the column metadata array, in the natural column-name order of [#getAllColumns()], and + /// empty for a CONSUMING segment. Like [#getAllColumns()] it is a snapshot. + @Override + public Collection<ColumnMetadata> getAllColumnMetadata() { + ColumnMetadata[] columnMetadata = _columnMetadata; + return columnMetadata != null ? Collections.unmodifiableList(Arrays.asList(columnMetadata)) : List.of(); + } + + @Override + public void forEachColumn(BiConsumer<String, ColumnMetadata> action) { + String[] columnNames = _columnNames; + if (columnNames == null) { + return; + } + ColumnMetadata[] columnMetadata = _columnMetadata; + for (int i = 0; i < columnNames.length; i++) { + action.accept(columnNames[i], columnMetadata[i]); + } + } + + @Nullable + @Override + public ColumnMetadata getColumnMetadataFor(String column) { + int index = indexOf(column); + return index >= 0 ? _columnMetadata[index] : null; + } + + /// Index of the column in the two arrays, or `-(insertion point) - 1`; always negative when there is no column + /// metadata at all (CONSUMING segment). + private int indexOf(String column) { + String[] columnNames = _columnNames; + return columnNames != null ? Arrays.binarySearch(columnNames, column) : -1; } @Override @@ -550,17 +615,100 @@ public class SegmentMetadataImpl implements SegmentMetadata { return _endOffset; } + /// {@inheritDoc} + /// + /// Built from the column arrays on the first call and cached until the columns change, so a caller pays one map + /// entry per column and the segment keeps it for its lifetime. Nothing on the load or query path should call this + /// — see the accessors listed on [SegmentMetadata#getColumnMetadataMap()]. Writes to the returned map do not reach + /// the segment metadata; use [#addColumnMetadata(String, ColumnMetadata)] and [#removeColumn(String)] instead. + /// + /// Returns `null` for a CONSUMING segment, which holds no column metadata. + @Nullable @Override public TreeMap<String, ColumnMetadata> getColumnMetadataMap() { - return _columnMetadataMap; + if (_columnNames == null) { + return null; + } + TreeMap<String, ColumnMetadata> columnMetadataMap = _columnMetadataMapView; + if (columnMetadataMap == null) { + synchronized (this) { + columnMetadataMap = _columnMetadataMapView; + if (columnMetadataMap == null) { + columnMetadataMap = buildColumnMetadataMap(); + _columnMetadataMapView = columnMetadataMap; + } + } + } + return columnMetadataMap; + } + + private TreeMap<String, ColumnMetadata> buildColumnMetadataMap() { + NUM_COLUMN_METADATA_MAP_MATERIALIZATIONS.incrementAndGet(); + TreeMap<String, ColumnMetadata> columnMetadataMap = new TreeMap<>(); + forEachColumn(columnMetadataMap::put); + return columnMetadataMap; + } + + /// Whether [#getColumnMetadataMap()] has been called (and its map cached) since the columns last changed. + @VisibleForTesting + public boolean isColumnMetadataMapMaterialized() { + return _columnMetadataMapView != null; + } + + /// Number of column metadata maps derived from the column arrays so far in this JVM. A load or query path that + /// leaves this unchanged did not build any segment's map. + @VisibleForTesting + public static long getNumColumnMetadataMapMaterializations() { + return NUM_COLUMN_METADATA_MAP_MATERIALIZATIONS.get(); + } + + /// {@inheritDoc} + /// + /// Inserts the column in natural order, which is a copy of both arrays; the loader adds a handful of virtual + /// columns once per segment, so this is not a hot path. + @Override + public void addColumnMetadata(String column, ColumnMetadata columnMetadata) { + String[] columnNames = _columnNames; + Preconditions.checkState(columnNames != null, "Segment: %s holds no column metadata", _segmentName); + int index = Arrays.binarySearch(columnNames, column); + if (index >= 0) { + _columnMetadata[index] = columnMetadata; + } else { + int insertionPoint = -index - 1; + _columnMetadata = insert(_columnMetadata, insertionPoint, columnMetadata); + _columnNames = insert(columnNames, insertionPoint, column); + } + invalidateDerivedViews(); } @Override public void removeColumn(String column) { Preconditions.checkState(!column.equals(_timeColumn), "Cannot remove time column: %s", _timeColumn); - _columnMetadataMap.remove(column); - // Drop the derived schema, if one was built, so the next getSchema() rebuilds it without the column + int index = indexOf(column); + if (index >= 0) { + _columnMetadata = delete(_columnMetadata, index); + _columnNames = delete(_columnNames, index); + } + invalidateDerivedViews(); + } + + /// Drops the schema and the map derived from the columns, so the next caller rebuilds them from the current arrays. + private void invalidateDerivedViews() { _schema = null; + _columnMetadataMapView = null; + } + + private static <E> E[] insert(E[] array, int index, E element) { + E[] extended = Arrays.copyOf(array, array.length + 1); + System.arraycopy(array, index, extended, index + 1, array.length - index); + extended[index] = element; + return extended; + } + + private static <E> E[] delete(E[] array, int index) { + E[] shortened = Arrays.copyOf(array, array.length - 1); + System.arraycopy(array, index + 1, shortened, index, array.length - index - 1); + return shortened; } @Override @@ -608,13 +756,13 @@ public class SegmentMetadataImpl implements SegmentMetadata { segmentMetadata.put("startOffset", _startOffset); segmentMetadata.put("endOffset", _endOffset); - if (_columnMetadataMap != null) { + if (_columnNames != null) { ArrayNode columnsMetadata = JsonUtils.newArrayNode(); - for (Map.Entry<String, ColumnMetadata> entry : _columnMetadataMap.entrySet()) { - if (columnFilter == null || columnFilter.contains(entry.getKey())) { - columnsMetadata.add(JsonUtils.objectToJsonNode(entry.getValue())); + forEachColumn((column, columnMetadata) -> { + if (columnFilter == null || columnFilter.contains(column)) { + columnsMetadata.add(JsonUtils.objectToJsonNode(columnMetadata)); } - } + }); segmentMetadata.set("columns", columnsMetadata); } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java new file mode 100644 index 00000000000..d8be4387e2d --- /dev/null +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java @@ -0,0 +1,264 @@ +/** + * 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.spi.index.metadata; + +import java.util.AbstractSet; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.Comparator; +import java.util.Iterator; +import java.util.NavigableSet; +import java.util.NoSuchElementException; +import java.util.SortedSet; +import java.util.TreeSet; +import javax.annotation.Nullable; + +import static com.google.common.base.Preconditions.checkArgument; + + +/// Unmodifiable [NavigableSet] view of a range of a sorted, duplicate-free `String[]`, ordered naturally. +/// +/// Lookups are a binary search over the array, so the whole set costs one small object rather than a red-black-tree +/// node per element. That is the point: a server holds the column names of every loaded segment for the segment's +/// lifetime, and a wide segment has thousands of them. +/// +/// The array is referenced, not copied, so the view reflects nothing the holder does afterwards *except* in-place +/// writes: [SegmentMetadataImpl] replaces its arrays when its columns change, which leaves an already-returned view +/// as the snapshot taken at the time of the call. +/// +/// Immutable and thread-safe as long as the backing array is not written in place. +final class SortedStringArraySet extends AbstractSet<String> implements NavigableSet<String> { + private final String[] _elements; + private final int _from; + private final int _to; + + SortedStringArraySet(String[] elements) { + this(elements, 0, elements.length); + } + + private SortedStringArraySet(String[] elements, int from, int to) { + _elements = elements; + _from = from; + _to = to; + } + + @Override + public int size() { + return _to - _from; + } + + @Override + public boolean isEmpty() { + return _from == _to; + } + + @Override + public boolean contains(Object o) { + return o instanceof String && search((String) o) >= 0; + } + + /// Index of `element`, or `-(insertion point) - 1`, both absolute in [#_elements]. + private int search(String element) { + return Arrays.binarySearch(_elements, _from, _to, element); + } + + private int ceilingIndex(String element) { + int index = search(element); + return index >= 0 ? index : -index - 1; + } + + private int higherIndex(String element) { + int index = search(element); + return index >= 0 ? index + 1 : -index - 1; + } + + private int floorIndex(String element) { + int index = search(element); + return index >= 0 ? index : -index - 2; + } + + private int lowerIndex(String element) { + int index = search(element); + return index >= 0 ? index - 1 : -index - 2; + } + + @Nullable + private String at(int index) { + return index >= _from && index < _to ? _elements[index] : null; + } + + @Nullable + @Override + public String ceiling(String element) { + return at(ceilingIndex(element)); + } + + @Nullable + @Override + public String higher(String element) { + return at(higherIndex(element)); + } + + @Nullable + @Override + public String floor(String element) { + return at(floorIndex(element)); + } + + @Nullable + @Override + public String lower(String element) { + return at(lowerIndex(element)); + } + + @Override + public String first() { + if (isEmpty()) { + throw new NoSuchElementException(); + } + return _elements[_from]; + } + + @Override + public String last() { + if (isEmpty()) { + throw new NoSuchElementException(); + } + return _elements[_to - 1]; + } + + @Nullable + @Override + public Comparator<? super String> comparator() { + return null; + } + + @Override + public Iterator<String> iterator() { + return new Iterator<>() { + private int _index = _from; + + @Override + public boolean hasNext() { + return _index < _to; + } + + @Override + public String next() { + if (_index >= _to) { + throw new NoSuchElementException(); + } + return _elements[_index++]; + } + }; + } + + @Override + public Iterator<String> descendingIterator() { + return new Iterator<>() { + private int _index = _to; + + @Override + public boolean hasNext() { + return _index > _from; + } + + @Override + public String next() { + if (_index <= _from) { + throw new NoSuchElementException(); + } + return _elements[--_index]; + } + }; + } + + @Override + public NavigableSet<String> subSet(String from, boolean fromInclusive, String to, boolean toInclusive) { + checkArgument(from.compareTo(to) <= 0, "from: %s > to: %s", from, to); + int start = fromInclusive ? ceilingIndex(from) : higherIndex(from); + int end = toInclusive ? higherIndex(to) : ceilingIndex(to); + return new SortedStringArraySet(_elements, start, Math.max(start, end)); + } + + @Override + public SortedSet<String> subSet(String from, String to) { + return subSet(from, true, to, false); + } + + @Override + public NavigableSet<String> headSet(String to, boolean inclusive) { + return new SortedStringArraySet(_elements, _from, inclusive ? higherIndex(to) : ceilingIndex(to)); + } + + @Override + public SortedSet<String> headSet(String to) { + return headSet(to, false); + } + + @Override + public NavigableSet<String> tailSet(String from, boolean inclusive) { + return new SortedStringArraySet(_elements, inclusive ? ceilingIndex(from) : higherIndex(from), _to); + } + + @Override + public SortedSet<String> tailSet(String from) { + return tailSet(from, true); + } + + /// Unlike the range sets above this is a copy, not a view, since the backing array is ascending. Nothing on the + /// segment paths calls it; it exists so the [NavigableSet] contract holds. + @Override + public NavigableSet<String> descendingSet() { + TreeSet<String> descending = new TreeSet<>(Comparator.reverseOrder()); + descending.addAll(this); + return Collections.unmodifiableNavigableSet(descending); + } + + @Override + public boolean remove(Object o) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean removeAll(Collection<?> c) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean retainAll(Collection<?> c) { + throw new UnsupportedOperationException(); + } + + @Override + public void clear() { + throw new UnsupportedOperationException(); + } + + @Override + public String pollFirst() { + throw new UnsupportedOperationException(); + } + + @Override + public String pollLast() { + throw new UnsupportedOperationException(); + } +} diff --git a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java new file mode 100644 index 00000000000..3564bc9500c --- /dev/null +++ b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java @@ -0,0 +1,106 @@ +/** + * 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.spi.index.metadata; + +import java.util.ArrayList; +import java.util.List; +import java.util.NavigableSet; +import java.util.NoSuchElementException; +import java.util.TreeSet; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/// The sorted-array view stands in for the `TreeMap` key set the segment metadata used to hand out, so it has to +/// answer every navigation query the same way a `TreeSet` of the same elements does. +public class SortedStringArraySetTest { + private static final String[] ELEMENTS = {"b", "d", "f", "h"}; + + private static NavigableSet<String> set() { + return new SortedStringArraySet(ELEMENTS.clone()); + } + + private static NavigableSet<String> reference() { + return new TreeSet<>(List.of(ELEMENTS)); + } + + @Test + public void testMatchesTreeSetForEveryProbe() { + NavigableSet<String> set = set(); + NavigableSet<String> reference = reference(); + assertEquals(set, reference); + assertEquals(set.hashCode(), reference.hashCode()); + assertEquals(set.size(), reference.size()); + assertEquals(new ArrayList<>(set), new ArrayList<>(reference)); + assertEquals(set.first(), reference.first()); + assertEquals(set.last(), reference.last()); + assertNull(set.comparator()); + for (String probe : List.of("a", "b", "c", "d", "e", "f", "g", "h", "i")) { + assertEquals(set.contains(probe), reference.contains(probe), probe); + assertEquals(set.floor(probe), reference.floor(probe), probe); + assertEquals(set.ceiling(probe), reference.ceiling(probe), probe); + assertEquals(set.lower(probe), reference.lower(probe), probe); + assertEquals(set.higher(probe), reference.higher(probe), probe); + assertEquals(new ArrayList<>(set.headSet(probe, true)), new ArrayList<>(reference.headSet(probe, true)), probe); + assertEquals(new ArrayList<>(set.headSet(probe, false)), new ArrayList<>(reference.headSet(probe, false)), probe); + assertEquals(new ArrayList<>(set.tailSet(probe, true)), new ArrayList<>(reference.tailSet(probe, true)), probe); + assertEquals(new ArrayList<>(set.tailSet(probe, false)), new ArrayList<>(reference.tailSet(probe, false)), probe); + assertEquals(new ArrayList<>(set.subSet(probe, true, "i", false)), + new ArrayList<>(reference.subSet(probe, true, "i", false)), probe); + } + assertEquals(new ArrayList<>(set.descendingSet()), new ArrayList<>(reference.descendingSet())); + List<String> descending = new ArrayList<>(); + set.descendingIterator().forEachRemaining(descending::add); + assertEquals(descending, List.of("h", "f", "d", "b")); + } + + @Test + public void testEmptyAndExhaustedIteration() { + NavigableSet<String> empty = new SortedStringArraySet(new String[0]); + assertTrue(empty.isEmpty()); + assertEquals(empty.size(), 0); + assertFalse(empty.iterator().hasNext()); + assertThrows(NoSuchElementException.class, empty::first); + assertThrows(NoSuchElementException.class, empty::last); + assertThrows(NoSuchElementException.class, () -> empty.iterator().next()); + assertThrows(NoSuchElementException.class, () -> empty.descendingIterator().next()); + assertFalse(set().subSet("c", true, "c", true).iterator().hasNext()); + } + + /// It is a view of the segment metadata's own array, so every mutator has to bounce rather than silently narrow + /// the columns of a loaded segment. + @Test + public void testUnmodifiable() { + NavigableSet<String> set = set(); + assertThrows(UnsupportedOperationException.class, () -> set.add("a")); + assertThrows(UnsupportedOperationException.class, () -> set.remove("b")); + assertThrows(UnsupportedOperationException.class, () -> set.removeAll(List.of("b"))); + assertThrows(UnsupportedOperationException.class, () -> set.retainAll(List.of("b"))); + assertThrows(UnsupportedOperationException.class, set::clear); + assertThrows(UnsupportedOperationException.class, set::pollFirst); + assertThrows(UnsupportedOperationException.class, set::pollLast); + assertThrows(UnsupportedOperationException.class, () -> set.iterator().remove()); + assertThrows(IllegalArgumentException.class, () -> set.subSet("f", true, "b", true)); + } +} diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReader.java b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReader.java index 7084843bc4c..cd2c959d66e 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReader.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReader.java @@ -80,7 +80,7 @@ final class SegmentCompressionStatsReader { includeColumnCompressionStats ? new HashMap<>() : null; IndexService indexService = includeColumnCompressionStats ? IndexService.getInstance() : null; - for (ColumnMetadata columnMetadata : segmentMetadata.getColumnMetadataMap().values()) { + for (ColumnMetadata columnMetadata : segmentMetadata.getAllColumnMetadata()) { long forwardIndexSize = getIndexSize(segmentMetadata, columnMetadata, StandardIndexes.forward()); if (forwardIndexSize < 0) { continue; diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java index d860f9c37df..38c3f5271c5 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java @@ -245,7 +245,9 @@ public class TablesResource { Set<String> allSegmentColumns = segmentMetadata.getAllColumns(); if (columnSet == null) { - columnSet = allSegmentColumns; + // Copy: getAllColumns() is an unmodifiable view of the segment's own columns, and retainAll below + // would otherwise narrow the first segment's metadata rather than the running intersection. + columnSet = new HashSet<>(allSegmentColumns); } else { columnSet.retainAll(allSegmentColumns); } @@ -253,7 +255,7 @@ public class TablesResource { // Column stats are scoped to the caller's column filter. for (String column : columnSet) { - ColumnMetadata columnMetadata = segmentMetadata.getColumnMetadataMap().get(column); + ColumnMetadata columnMetadata = segmentMetadata.getColumnMetadataFor(column); int columnLength = columnMetadata.getLengthOfLongestElement(); if (columnLength < 0) { // For raw STRING/BYTES/BIG_DECIMAL column, set the columnLength as the length of the max value. @@ -385,10 +387,10 @@ public class TablesResource { @Nullable Set<String> columnFilter) { int additionalCount = 0; if (columnFilter == null) { - additionalCount = segment.getSegmentMetadata().getColumnMetadataMap().size(); + additionalCount = segment.getSegmentMetadata().getNumColumns(); } else { for (String column : columnFilter) { - if (segment.getSegmentMetadata().getColumnMetadataMap().containsKey(column)) { + if (segment.getSegmentMetadata().getColumnMetadataFor(column) != null) { additionalCount++; } } diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReaderTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReaderTest.java index 91720990279..69a3195757a 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReaderTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/resources/SegmentCompressionStatsReaderTest.java @@ -22,7 +22,6 @@ import java.io.File; import java.nio.file.Files; import java.util.ArrayList; import java.util.List; -import java.util.Map; import java.util.TreeMap; import javax.ws.rs.WebApplicationException; import org.apache.commons.io.FileUtils; @@ -61,7 +60,7 @@ public class SegmentCompressionStatsReaderTest { public void testEmptySegmentHasCompleteZeroByteStats() { SegmentMetadata segmentMetadata = mock(SegmentMetadata.class); when(segmentMetadata.getName()).thenReturn("empty"); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>()); + when(segmentMetadata.getAllColumnMetadata()).thenReturn(List.of()); assertCompleteZeroByteStats(SegmentCompressionStatsReader.read(segmentMetadata, true)); } @@ -74,7 +73,7 @@ public class SegmentCompressionStatsReaderTest { when(segmentMetadata.getName()).thenReturn("forwardDisabled"); TreeMap<String, ColumnMetadata> columnMetadataMap = new TreeMap<>(); columnMetadataMap.put("column", columnMetadata); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + when(segmentMetadata.getAllColumnMetadata()).thenReturn(columnMetadataMap.values()); assertCompleteZeroByteStats(SegmentCompressionStatsReader.read(segmentMetadata, false)); } @@ -123,8 +122,7 @@ public class SegmentCompressionStatsReaderTest { @Test public void testServerRejectsOversizedColumnContributionResponse() { SegmentMetadata segmentMetadata = mock(SegmentMetadata.class); - when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>(Map.of( - "column", mock(ColumnMetadata.class)))); + when(segmentMetadata.getNumColumns()).thenReturn(1); ImmutableSegment segment = mock(ImmutableSegment.class); when(segment.getSegmentMetadata()).thenReturn(segmentMetadata); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
