From 187be1c87b4884677b1d2caffca0910f108b89c1 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 6 Sep 2026 16:05:58 -0700 Subject: [PATCH 1/2] DATA-3221 (11): list physical columns without building the segment schema Two things the preprocess does on every segment load asked the segment metadata for its schema, and since the schema is derived and then cached, each one pinned a per-segment `Schema` for the segment's whole life: - `ForwardIndexHandler#computeOperations` needs the set of physical column names, and - `ColumnMinMaxValueGenerator` needs the columns its mode selects (the default mode is `ALL`, so this runs on every load). Both questions are answered by the column metadata the schema is itself derived from. `SegmentMetadata#getPhysicalColumnNames()` walks the column metadata for the first, falling back to the schema for a segment that holds no column metadata (a CONSUMING one), and the min/max generator now selects straight off each column's field spec. On a server measured with 13.6k loaded segments, the schemas built here were ~144 MB of tree entries and list slots, all of it a second copy of data the column metadata already holds. Co-Authored-By: Claude Opus 5 --- .../index/loader/ForwardIndexHandler.java | 2 +- .../ColumnMinMaxValueGenerator.java | 70 ++++++++----------- .../index/SegmentMetadataImplTest.java | 35 ++++++++++ .../pinot/segment/spi/SegmentMetadata.java | 22 ++++++ 4 files changed, 86 insertions(+), 43 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java index 9c53152cdee3..d9a5a93fa502 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java @@ -248,7 +248,7 @@ Map> computeOperations(SegmentDirectory.Reader segmentRe } Map> columnOperationsMap = new HashMap<>(); - Set existingAllColumns = segmentMetadata.getSchema().getPhysicalColumnNames(); + Set existingAllColumns = segmentMetadata.getPhysicalColumnNames(); Set existingDictColumns = _segmentDirectory.getColumnsWithIndex(StandardIndexes.dictionary()); Set existingForwardIndexColumns = _segmentDirectory.getColumnsWithIndex(StandardIndexes.forward()); Set existingInvertedIndexColumns = diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java index ae4b99f84790..3ea90dc46cf1 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java @@ -23,6 +23,7 @@ import java.math.BigDecimal; import java.util.ArrayList; import java.util.List; +import java.util.function.Consumer; import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.commons.lang3.Strings; import org.apache.pinot.segment.local.segment.creator.impl.SegmentColumnarIndexCreator; @@ -44,7 +45,6 @@ import org.apache.pinot.segment.spi.store.SegmentDirectory; import org.apache.pinot.segment.spi.utils.SegmentMetadataUtils; import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.ByteArray; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -75,11 +75,11 @@ public ColumnMinMaxValueGenerator(SegmentMetadata segmentMetadata, SegmentDirect /// Returns the list of columns that need min/max values to be updated public List columnMinMaxValueUpdates() { List columns = new ArrayList<>(); - for (String column : getColumnsToAddMinMaxValue()) { - if (needAddColumnMinMaxValueForColumn(column)) { - columns.add(column); + forEachSelectedColumn(columnMetadata -> { + if (needAddColumnMinMaxValueForColumn(columnMetadata)) { + columns.add(columnMetadata.getColumnName()); } - } + }); return columns; } @@ -87,53 +87,40 @@ public void addColumnMinMaxValue() throws Exception { Preconditions.checkState(_columnMinMaxValueGeneratorMode != ColumnMinMaxValueGeneratorMode.NONE); _segmentProperties = SegmentMetadataUtils.getPropertiesConfiguration(_segmentMetadata); - for (String column : getColumnsToAddMinMaxValue()) { - addColumnMinMaxValueForColumn(column); - } + forEachSelectedColumn(this::addColumnMinMaxValueForColumn); if (_minMaxValueAdded) { SegmentMetadataUtils.savePropertiesConfiguration(_segmentProperties, _segmentMetadata.getIndexDir()); } } - private List getColumnsToAddMinMaxValue() { - Schema schema = _segmentMetadata.getSchema(); - List columnsToAddMinMaxValue = new ArrayList<>(); + /// Runs `action` on every column the generator mode selects. + /// + /// The selection reads the field specs off the column metadata rather than off `_segmentMetadata.getSchema()`, + /// which is the same data (the schema is derived from the column metadata) but costs a `Schema` per segment. This + /// runs on every segment load — the default mode is `ALL` — so a schema built here would be cached for the + /// segment's whole life, and a server holding tens of thousands of wide segments would keep one per segment. + private void forEachSelectedColumn(Consumer action) { + for (ColumnMetadata columnMetadata : _segmentMetadata.getAllColumnMetadata()) { + FieldSpec fieldSpec = columnMetadata.getFieldSpec(); + if (!fieldSpec.isVirtualColumn() && isSelected(fieldSpec.getFieldType())) { + action.accept(columnMetadata); + } + } + } - // mode ALL - use all columns - // mode NON_METRIC - use all dimensions and time columns - // mode TIME - use only time columns + /// Whether the generator mode covers the given field type: `ALL` takes every column, `NON_METRIC` every column but + /// the metrics, `TIME` only the time columns. + private boolean isSelected(FieldSpec.FieldType fieldType) { switch (_columnMinMaxValueGeneratorMode) { case ALL: - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn()) { - columnsToAddMinMaxValue.add(fieldSpec.getName()); - } - } - break; + return true; case NON_METRIC: - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn() && fieldSpec.getFieldType() != FieldSpec.FieldType.METRIC) { - columnsToAddMinMaxValue.add(fieldSpec.getName()); - } - } - break; + return fieldType != FieldSpec.FieldType.METRIC; case TIME: - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn() && (fieldSpec.getFieldType() == FieldSpec.FieldType.TIME - || fieldSpec.getFieldType() == FieldSpec.FieldType.DATE_TIME)) { - columnsToAddMinMaxValue.add(fieldSpec.getName()); - } - } - break; + return fieldType == FieldSpec.FieldType.TIME || fieldType == FieldSpec.FieldType.DATE_TIME; default: throw new IllegalStateException("Unsupported generator mode: " + _columnMinMaxValueGeneratorMode); } - - return columnsToAddMinMaxValue; - } - - private boolean needAddColumnMinMaxValueForColumn(String columnName) { - return needAddColumnMinMaxValueForColumn(_segmentMetadata.getColumnMetadataFor(columnName)); } private boolean needAddColumnMinMaxValueForColumn(ColumnMetadata columnMetadata) { @@ -141,8 +128,7 @@ private boolean needAddColumnMinMaxValueForColumn(ColumnMetadata columnMetadata) && !columnMetadata.isMinMaxValueInvalid(); } - private void addColumnMinMaxValueForColumn(String columnName) { - ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(columnName); + private void addColumnMinMaxValueForColumn(ColumnMetadata columnMetadata) { if (!needAddColumnMinMaxValueForColumn(columnMetadata)) { return; } @@ -155,7 +141,7 @@ private void addColumnMinMaxValueForColumn(String columnName) { _minMaxValueAdded = true; } catch (Exception e) { LOGGER.error("Caught exception while generating min/max value for column: {} in segment: {}, continuing without " - + "persisting them", columnName, _segmentMetadata.getName(), e); + + "persisting them", columnMetadata.getColumnName(), _segmentMetadata.getName(), e); } } 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 31ede949bbea..af97763a2bef 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 @@ -40,7 +40,10 @@ import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.converter.SegmentV1V2ToV3FormatConverter; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.local.segment.index.loader.SegmentPreProcessor; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.store.SegmentLocalFSDirectory; import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProviderFactory; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.ImmutableSegment; @@ -51,6 +54,7 @@ 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.SegmentDirectory; import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.config.table.OpenStructIndexConfig; @@ -332,6 +336,37 @@ public void testSchemaDerivedLazilyFromColumnMetadata() assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), materializations + 1); } + /// The preprocess that runs on every segment load asks the forward-index handler which physical columns exist. That + /// question must not build the per-segment schema: doing so once per segment pins one [Schema] per loaded segment + /// for its whole life, which on a server holding tens of thousands of wide segments is hundreds of megabytes. + @Test + public void testPreprocessDoesNotBuildTheSegmentSchema() + throws Exception { + // The forward-index handler skips segments older than v3, so the preprocess only reaches it on a v3 segment. + new SegmentV1V2ToV3FormatConverter().convert(_segmentDirectory); + + long materializations = SegmentMetadataImpl.getNumSchemaMaterializations(); + SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory); + Set physical = metadata.getPhysicalColumnNames(); + assertFalse(metadata.isSchemaMaterialized(), "listing physical columns must not build the segment schema"); + assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), materializations); + assertEquals(physical, metadata.getSchema().getPhysicalColumnNames(), + "the derived names must equal what the schema reports"); + assertFalse(physical.contains(BuiltInVirtualColumn.DOCID)); + + TableConfig tableConfig = + new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable").setTimeColumnName("daysSinceEpoch").build(); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(tableConfig, metadata.getSchema()); + indexLoadingConfig.setReadMode(ReadMode.mmap); + long beforePreprocess = SegmentMetadataImpl.getNumSchemaMaterializations(); + try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(_segmentDirectory, ReadMode.mmap); + SegmentPreProcessor preProcessor = new SegmentPreProcessor(segmentDirectory, indexLoadingConfig)) { + preProcessor.process(); + } + assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), beforePreprocess, + "segment preprocess must not build any segment schema"); + } + /// Loading a segment registers the built-in virtual columns in the column metadata, so the schema derived afterwards /// includes them exactly as the schema the loader used to build eagerly did, while neither the load nor serving the /// segment (column listings, data sources, the metadata JSON) builds any schema. 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 142bf688dcaf..a2caae4105ea 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 @@ -19,6 +19,7 @@ package org.apache.pinot.segment.spi; import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.collect.Sets; import java.io.File; import java.util.Collection; import java.util.List; @@ -33,6 +34,7 @@ import org.apache.pinot.segment.spi.index.multicolumntext.MultiColumnTextMetadata; import org.apache.pinot.segment.spi.index.startree.StarTreeV2Metadata; import org.apache.pinot.spi.annotations.InterfaceAudience; +import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.joda.time.Duration; import org.joda.time.Interval; @@ -161,6 +163,26 @@ default ColumnMetadata getColumnMetadataFor(String column) { return getColumnMetadataMap().get(column); } + /// The names of the physical (non-virtual) columns, i.e. `getSchema().getPhysicalColumnNames()` without building + /// the schema. Segment load runs this once per segment (the forward-index handler asks which columns exist), and on + /// a server holding tens of thousands of wide segments a schema built there would be cached for the segment's whole + /// life: one [Schema] per segment, each with a tree entry and two list slots per column. A segment that holds no + /// column metadata (a CONSUMING one) still answers from its schema, which it was constructed with. + default Set getPhysicalColumnNames() { + Collection columnMetadata = getAllColumnMetadata(); + if (columnMetadata.isEmpty()) { + return getSchema().getPhysicalColumnNames(); + } + Set physicalColumnNames = Sets.newHashSetWithExpectedSize(columnMetadata.size()); + for (ColumnMetadata metadata : columnMetadata) { + FieldSpec fieldSpec = metadata.getFieldSpec(); + if (!fieldSpec.isVirtualColumn()) { + physicalColumnNames.add(fieldSpec.getName()); + } + } + return physicalColumnNames; + } + /// Registers the metadata of a column, replacing any metadata already registered under the same name. An /// implementation that holds no column metadata (a CONSUMING segment) may reject this. default void addColumnMetadata(String column, ColumnMetadata columnMetadata) { From 1bad6fdfc20b0f234d9864def26433ac032d380c Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 6 Sep 2026 16:58:32 -0700 Subject: [PATCH 2/2] DATA-3221 (12): make the new column-name accessor sorted and null-safe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two follow-ups on `SegmentMetadata#getPhysicalColumnNames()`, both found by review of the previous commit rather than by a failure: - It returned a `HashSet` for a segment that holds column metadata but the schema's `TreeSet` for one that does not, so the expression it replaced (`getSchema().getPhysicalColumnNames()`, always sorted) silently became unordered for immutable segments and stayed sorted for CONSUMING ones — the worst shape for a caller that assumes order. It now returns a `SortedSet` for both, as the replaced expression did. - Its schema fallback for a segment holding no column metadata was reached through the default `getAllColumnMetadata()`, which is `getColumnMetadataMap().values()` — and `getColumnMetadataMap()` is `@Nullable`, documented to answer `null` for exactly that segment. The fallback worked only because the one implementation in the tree overrides the accessor. That default now honors its own documented contract and answers empty for a null map. Co-Authored-By: Claude Opus 5 --- .../index/SegmentMetadataImplTest.java | 4 +- .../pinot/segment/spi/SegmentMetadata.java | 12 ++- .../spi/SegmentMetadataDefaultsTest.java | 87 +++++++++++++++++++ 3 files changed, 97 insertions(+), 6 deletions(-) create mode 100644 pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/SegmentMetadataDefaultsTest.java 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 af97763a2bef..9870849ba334 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 @@ -350,8 +350,8 @@ public void testPreprocessDoesNotBuildTheSegmentSchema() Set physical = metadata.getPhysicalColumnNames(); assertFalse(metadata.isSchemaMaterialized(), "listing physical columns must not build the segment schema"); assertEquals(SegmentMetadataImpl.getNumSchemaMaterializations(), materializations); - assertEquals(physical, metadata.getSchema().getPhysicalColumnNames(), - "the derived names must equal what the schema reports"); + assertEquals(List.copyOf(physical), List.copyOf(metadata.getSchema().getPhysicalColumnNames()), + "the derived names must equal what the schema reports, in the same order"); assertFalse(physical.contains(BuiltInVirtualColumn.DOCID)); TableConfig tableConfig = 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 a2caae4105ea..5e8ddaca2784 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 @@ -19,14 +19,15 @@ package org.apache.pinot.segment.spi; import com.fasterxml.jackson.databind.JsonNode; -import com.google.common.collect.Sets; 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.SortedSet; import java.util.TreeMap; +import java.util.TreeSet; import java.util.concurrent.TimeUnit; import java.util.function.BiConsumer; import javax.annotation.Nullable; @@ -136,7 +137,8 @@ default int getNumColumns() { /// empty for a segment that holds none (a CONSUMING one, which answers [#getColumnMetadataFor(String)] with `null` /// for every column of its schema). default Collection getAllColumnMetadata() { - return getColumnMetadataMap().values(); + TreeMap columnMetadataMap = getColumnMetadataMap(); + return columnMetadataMap != null ? columnMetadataMap.values() : List.of(); } /// Applies `action` to every (column name, column metadata) pair, in the natural column-name order of @@ -168,12 +170,14 @@ default ColumnMetadata getColumnMetadataFor(String column) { /// a server holding tens of thousands of wide segments a schema built there would be cached for the segment's whole /// life: one [Schema] per segment, each with a tree entry and two list slots per column. A segment that holds no /// column metadata (a CONSUMING one) still answers from its schema, which it was constructed with. - default Set getPhysicalColumnNames() { + /// + /// Sorted, like the [Schema#getPhysicalColumnNames()] this replaces, and the same set for both segment kinds. + default SortedSet getPhysicalColumnNames() { Collection columnMetadata = getAllColumnMetadata(); if (columnMetadata.isEmpty()) { return getSchema().getPhysicalColumnNames(); } - Set physicalColumnNames = Sets.newHashSetWithExpectedSize(columnMetadata.size()); + TreeSet physicalColumnNames = new TreeSet<>(); for (ColumnMetadata metadata : columnMetadata) { FieldSpec fieldSpec = metadata.getFieldSpec(); if (!fieldSpec.isVirtualColumn()) { diff --git a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/SegmentMetadataDefaultsTest.java b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/SegmentMetadataDefaultsTest.java new file mode 100644 index 000000000000..753426a432bd --- /dev/null +++ b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/SegmentMetadataDefaultsTest.java @@ -0,0 +1,87 @@ +/** + * 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; + +import java.util.List; +import java.util.SortedSet; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; +import static org.testng.Assert.assertEquals; + + +/// Covers the [SegmentMetadata] default methods against an implementation that holds no column metadata, which +/// [SegmentMetadata#getColumnMetadataMap()] is documented to answer with `null`. The only implementation in the tree +/// overrides [SegmentMetadata#getAllColumnMetadata()], so the defaults are exercised here rather than through it. +public class SegmentMetadataDefaultsTest { + + @Test + public void testAllColumnMetadataIsEmptyWithoutAColumnMetadataMap() { + SegmentMetadata segmentMetadata = metadataHoldingNoColumns(); + assertEquals(segmentMetadata.getAllColumnMetadata(), List.of()); + } + + @Test + public void testPhysicalColumnNamesFallBackToTheSchemaWithoutColumnMetadata() { + SegmentMetadata segmentMetadata = metadataHoldingNoColumns(); + assertEquals(segmentMetadata.getPhysicalColumnNames(), schema().getPhysicalColumnNames()); + } + + @Test + public void testPhysicalColumnNamesAreSortedAndSkipVirtualColumns() { + ColumnMetadata zebra = columnMetadata("zebra", false); + ColumnMetadata apple = columnMetadata("apple", false); + ColumnMetadata virtual = columnMetadata("$docId", true); + SegmentMetadata segmentMetadata = mock(SegmentMetadata.class, CALLS_REAL_METHODS); + when(segmentMetadata.getAllColumnMetadata()).thenReturn(List.of(zebra, apple, virtual)); + + SortedSet physicalColumnNames = segmentMetadata.getPhysicalColumnNames(); + assertEquals(List.copyOf(physicalColumnNames), List.of("apple", "zebra")); + } + + private static SegmentMetadata metadataHoldingNoColumns() { + SegmentMetadata segmentMetadata = mock(SegmentMetadata.class, withSettings().defaultAnswer(CALLS_REAL_METHODS)); + when(segmentMetadata.getColumnMetadataMap()).thenReturn(null); + when(segmentMetadata.getSchema()).thenReturn(schema()); + return segmentMetadata; + } + + private static Schema schema() { + return new Schema.SchemaBuilder() + .addSingleValueDimension("zebra", FieldSpec.DataType.STRING) + .addSingleValueDimension("apple", FieldSpec.DataType.STRING) + .build(); + } + + private static ColumnMetadata columnMetadata(String column, boolean virtual) { + DimensionFieldSpec fieldSpec = new DimensionFieldSpec(column, FieldSpec.DataType.STRING, true); + if (virtual) { + fieldSpec.setVirtualColumnProvider("org.apache.pinot.segment.spi.virtualcolumn.DocIdVirtualColumnProvider"); + } + ColumnMetadata columnMetadata = mock(ColumnMetadata.class); + when(columnMetadata.getFieldSpec()).thenReturn(fieldSpec); + return columnMetadata; + } +}