This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 4eb02338125 Star-tree Index Builder: Support dimension columns with
raw forward index + shared dictionary (#18637)
4eb02338125 is described below
commit 4eb023381252729c11c41d9385873870b0a64e50
Author: Chaitanya Deepthi <[email protected]>
AuthorDate: Wed Jun 10 11:14:49 2026 -0700
Star-tree Index Builder: Support dimension columns with raw forward index
+ shared dictionary (#18637)
---
.../segment/creator/impl/BaseSegmentCreator.java | 9 +-
.../segment/index/loader/ForwardIndexHandler.java | 9 +-
.../segment/index/loader/SegmentPreProcessor.java | 3 +-
.../defaultcolumn/BaseDefaultColumnHandler.java | 4 +-
.../segment/readers/PinotSegmentColumnReader.java | 35 ++++-
.../segment/store/SingleFileIndexDirectory.java | 5 +
.../startree/v2/builder/MultipleTreesBuilder.java | 49 ++++++-
.../index/loader/ForwardIndexHandlerTest.java | 5 +-
...otSegmentColumnReaderRawWithDictionaryTest.java | 148 +++++++++++++++++++++
9 files changed, 250 insertions(+), 17 deletions(-)
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java
index 12d9cf88bdc..3938a6437b3 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java
@@ -865,10 +865,11 @@ public abstract class BaseSegmentCreator implements
SegmentCreator {
List<StarTreeIndexConfig> starTreeIndexConfigs =
_config.getStarTreeIndexConfigs();
boolean enableDefaultStarTree = _config.isEnableDefaultStarTree();
if (CollectionUtils.isNotEmpty(starTreeIndexConfigs) ||
enableDefaultStarTree) {
- MultipleTreesBuilder.BuildMode buildMode =
- _config.isOnHeap() ? MultipleTreesBuilder.BuildMode.ON_HEAP :
MultipleTreesBuilder.BuildMode.OFF_HEAP;
- MultipleTreesBuilder builder = new
MultipleTreesBuilder(starTreeIndexConfigs, enableDefaultStarTree, indexDir,
- buildMode);
+ // SegmentGeneratorConfig carries the star-tree configs, build mode, and
the TableConfig +
+ // Schema needed to construct an IndexLoadingConfig so downstream
readers (e.g. external-table
+ // forward-index readers backed by remote storage) can resolve
table-level configs during the
+ // in-place segment load.
+ MultipleTreesBuilder builder = new MultipleTreesBuilder(indexDir,
_config);
// We don't create the builder using the try-with-resources pattern
because builder.close() performs
// some clean-up steps to roll back the star-tree index to the previous
state if it exists. If this goes wrong
// the star-tree index can be in an inconsistent state. To prevent that,
when builder.close() throws an
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 e0f2e200fd0..1df5608e596 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
@@ -874,7 +874,8 @@ public class ForwardIndexHandler extends BaseIndexHandler {
SegmentDictionaryCreator dictionaryCreator) {
boolean isSVColumn = reader.isSingleValue();
int maxNumValuesPerEntry =
existingColumnMetadata.getMaxNumberOfMultiValues();
- PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(reader, null, null, maxNumValuesPerEntry);
+ PinotSegmentColumnReader columnReader =
+ new PinotSegmentColumnReader(column, reader, null, null,
maxNumValuesPerEntry);
for (int i = 0; i < numDocs; i++) {
Object obj = columnReader.getValue(i);
@@ -936,7 +937,7 @@ public class ForwardIndexHandler extends BaseIndexHandler {
// Special null handling is not necessary here. This is because, the
existing default null value in the raw
// forwardIndex will be retained as such while created the dictionary
and dict-based forward index. Also, null
// value vectors maintain a bitmap of docIds. No handling is necessary
there.
- try (PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(reader, null, null,
+ try (PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(column, reader, null, null,
existingColMetadata.getMaxNumberOfMultiValues())) {
for (int i = 0; i < numDocs; i++) {
statsCollector.collect(columnReader.getValue(i));
@@ -1017,7 +1018,7 @@ public class ForwardIndexHandler extends BaseIndexHandler
{
.createIndexReader(segmentWriter, _fieldIndexConfigs.get(column),
existingColMetadata)) {
AbstractColumnStatisticsCollector statsCollector =
getStatsCollector(column, fieldSpec.getDataType().getStoredType(),
true);
- try (PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(reader, null, null,
+ try (PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(column, reader, null, null,
existingColMetadata.getMaxNumberOfMultiValues())) {
for (int i = 0; i < existingColMetadata.getTotalDocs(); i++) {
statsCollector.collect(columnReader.getValue(i));
@@ -1212,7 +1213,7 @@ public class ForwardIndexHandler extends BaseIndexHandler
{
}
boolean dictionaryEncoded = forwardIndex.isDictionaryEncoded();
try (Dictionary dictionary = dictionaryEncoded ?
DictionaryIndexType.read(segmentWriter, columnMetadata) : null;
- PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(forwardIndex, dictionary, null,
+ PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(column, forwardIndex, dictionary, null,
columnMetadata.getMaxNumberOfMultiValues())) {
AbstractColumnStatisticsCollector statsCollector =
getStatsCollector(column, storedType, false);
int numDocs = columnMetadata.getTotalDocs();
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessor.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessor.java
index cef94ee26ba..9faf9179b82 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessor.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessor.java
@@ -398,8 +398,9 @@ public class SegmentPreProcessor implements AutoCloseable {
StarTreeBuilderUtils.removeStarTrees(indexDir);
} else {
// NOTE: Always use OFF_HEAP mode on server side.
+ // Pass _indexLoadingConfig so downstream readers can resolve
table-level configs we set
MultipleTreesBuilder builder = new
MultipleTreesBuilder(starTreeBuilderConfigs, indexDir,
- MultipleTreesBuilder.BuildMode.OFF_HEAP);
+ MultipleTreesBuilder.BuildMode.OFF_HEAP, _indexLoadingConfig);
// We don't create the builder using the try-with-resources pattern
because builder.close() performs
// some clean-up steps to roll back the star-tree index to the
previous state if it exists. If this goes wrong
// the star-tree index can be in an inconsistent state. To prevent
that, when builder.close() throws an
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 d76b6209ca4..d88b3b90c60 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
@@ -1236,8 +1236,8 @@ public abstract class BaseDefaultColumnHandler implements
DefaultColumnHandler {
} else {
_dictionary = null;
}
- _columnReader = new PinotSegmentColumnReader(_forwardIndexReader,
_dictionary, null,
- columnMetadata.getMaxNumberOfMultiValues());
+ _columnReader = new
PinotSegmentColumnReader(columnMetadata.getColumnName(), _forwardIndexReader,
_dictionary,
+ null, columnMetadata.getMaxNumberOfMultiValues());
}
Object getValue(int docId) {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
index e5ac0e94f75..39ee9a240cb 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
@@ -42,6 +42,8 @@ public class PinotSegmentColumnReader implements Closeable {
private final NullValueVectorReader _nullValueVectorReader;
private final int[] _dictIdBuffer;
private final DataType _valueType;
+ @Nullable
+ private final String _columnName;
public PinotSegmentColumnReader(IndexSegment indexSegment, String column) {
DataSource dataSource = indexSegment.getDataSource(column);
@@ -51,6 +53,7 @@ public class PinotSegmentColumnReader implements Closeable {
_dictionary = dataSource.getDictionary();
_nullValueVectorReader = dataSource.getNullValueVector();
_valueType = _dictionary != null ? _dictionary.getValueType() :
_forwardIndexReader.getStoredType();
+ _columnName = column;
if (_forwardIndexReader.isSingleValue()) {
_dictIdBuffer = null;
} else {
@@ -60,13 +63,14 @@ public class PinotSegmentColumnReader implements Closeable {
}
}
- public PinotSegmentColumnReader(ForwardIndexReader forwardIndexReader,
@Nullable Dictionary dictionary,
+ public PinotSegmentColumnReader(String column, ForwardIndexReader
forwardIndexReader, @Nullable Dictionary dictionary,
@Nullable NullValueVectorReader nullValueVectorReader, int
maxNumValuesPerMVEntry) {
_forwardIndexReader = forwardIndexReader;
_forwardIndexReaderContext = _forwardIndexReader.createContext();
_dictionary = dictionary;
_nullValueVectorReader = nullValueVectorReader;
_valueType = _dictionary != null ? _dictionary.getValueType() :
_forwardIndexReader.getStoredType();
+ _columnName = column;
if (_forwardIndexReader.isSingleValue()) {
_dictIdBuffer = null;
} else {
@@ -92,7 +96,34 @@ public class PinotSegmentColumnReader implements Closeable {
}
public int getDictId(int docId) {
- return _forwardIndexReader.getDictId(docId, _forwardIndexReaderContext);
+ if (_forwardIndexReader.isDictionaryEncoded()) {
+ return _forwardIndexReader.getDictId(docId, _forwardIndexReaderContext);
+ }
+ if (_dictionary == null) {
+ throw new UnsupportedOperationException(
+ "Cannot resolve dictId: forward index is raw and no dictionary is
materialized for column: " + (
+ _columnName != null ? _columnName : "<unknown>"));
+ }
+ // If we have separate dictionary on a RAW forward index column, use that
dictionary.
+ switch (_valueType.getStoredType()) {
+ case INT:
+ return _dictionary.indexOf(_forwardIndexReader.getInt(docId,
_forwardIndexReaderContext));
+ case LONG:
+ return _dictionary.indexOf(_forwardIndexReader.getLong(docId,
_forwardIndexReaderContext));
+ case FLOAT:
+ return _dictionary.indexOf(_forwardIndexReader.getFloat(docId,
_forwardIndexReaderContext));
+ case DOUBLE:
+ return _dictionary.indexOf(_forwardIndexReader.getDouble(docId,
_forwardIndexReaderContext));
+ case BIG_DECIMAL:
+ return _dictionary.indexOf(_forwardIndexReader.getBigDecimal(docId,
_forwardIndexReaderContext));
+ case STRING:
+ return _dictionary.indexOf(_forwardIndexReader.getString(docId,
_forwardIndexReaderContext));
+ case BYTES:
+ return _dictionary.indexOf(new
ByteArray(_forwardIndexReader.getBytes(docId, _forwardIndexReaderContext)));
+ default:
+ throw new UnsupportedOperationException(
+ "Cannot resolve dictId for raw forward index of stored type: " +
_valueType.getStoredType());
+ }
}
public Object getValue(int docId) {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectory.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectory.java
index 3f72d69eb87..19119bcc940 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectory.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectory.java
@@ -345,6 +345,11 @@ class SingleFileIndexDirectory extends
ColumnIndexDirectory {
properties.putAll(_segmentDirectoryLoaderContext.getSegmentCustomConfigs());
}
+ // Propagate segment-level custom metadata so downstream readers can read
it
+ if (_segmentMetadata != null && _segmentMetadata.getCustomMap() != null) {
+ properties.putAll(_segmentMetadata.getCustomMap());
+ }
+
// Propagate the table's task config (serialized as JSON) to remote/empty
index buffers so that
// downstream readers backed by external storage can resolve any
configuration they require
// (e.g. credentials, regions, endpoints) from the ingestion task config
rather than relying
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/MultipleTreesBuilder.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/MultipleTreesBuilder.java
index 8de13eef94c..2ec4e331bac 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/MultipleTreesBuilder.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/MultipleTreesBuilder.java
@@ -34,12 +34,14 @@ import
org.apache.commons.configuration2.PropertiesConfiguration;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.tuple.Pair;
import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
import org.apache.pinot.segment.local.startree.StarTreeBuilderUtils;
import org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexMapUtils;
import
org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexMapUtils.IndexKey;
import
org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexMapUtils.IndexValue;
import org.apache.pinot.segment.spi.ImmutableSegment;
import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
import org.apache.pinot.segment.spi.index.startree.StarTreeV2Constants;
import
org.apache.pinot.segment.spi.index.startree.StarTreeV2Constants.MetadataKey;
@@ -86,6 +88,12 @@ public class MultipleTreesBuilder implements Closeable {
*/
public MultipleTreesBuilder(List<StarTreeV2BuilderConfig> builderConfigs,
File indexDir, BuildMode buildMode)
throws Exception {
+ this(builderConfigs, indexDir, buildMode, null);
+ }
+
+ public MultipleTreesBuilder(List<StarTreeV2BuilderConfig> builderConfigs,
File indexDir, BuildMode buildMode,
+ @Nullable IndexLoadingConfig indexLoadingConfig)
+ throws Exception {
Preconditions.checkArgument(CollectionUtils.isNotEmpty(builderConfigs),
"Must provide star-tree builder configs");
_builderConfigs = builderConfigs;
_buildMode = buildMode;
@@ -103,10 +111,26 @@ public class MultipleTreesBuilder implements Closeable {
}
LOGGER.debug(logUpdatedStarTrees.toString());
}
- _segment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap);
+ _segment = loadSegment(indexDir, indexLoadingConfig);
_starTreeCreationFailed = false;
}
+ private static ImmutableSegment loadSegment(File indexDir, @Nullable
IndexLoadingConfig indexLoadingConfig)
+ throws Exception {
+ if (indexLoadingConfig != null && indexLoadingConfig.getTableConfig() !=
null) {
+ // Minimal IndexLoadingConfig carrying only TableConfig. Schema is
intentionally null so the
+ // segment loader uses the segment's own column set rather than
filtering against an in-flight
+ // schema (e.g. one that adds/removes columns during pre-processing).
Also do NOT propagate
+ // segmentTier/segmentDirectoryLoader from the parent — a tier loader
would treat this as a
+ // tier-resolution event and relocate the segment mid-build. TableConfig
threads table-level
+ // config to downstream readers (e.g. external-storage forward-index
readers).
+ IndexLoadingConfig localConfig = new
IndexLoadingConfig(indexLoadingConfig.getTableConfig(), null);
+ localConfig.setReadMode(ReadMode.mmap);
+ return ImmutableSegmentLoader.load(indexDir, localConfig, false);
+ }
+ return ImmutableSegmentLoader.load(indexDir, ReadMode.mmap);
+ }
+
/**
* Constructor for the multiple star-trees builder.
*
@@ -118,6 +142,27 @@ public class MultipleTreesBuilder implements Closeable {
public MultipleTreesBuilder(@Nullable List<StarTreeIndexConfig>
indexConfigs, boolean enableDefaultStarTree,
File indexDir, BuildMode buildMode)
throws Exception {
+ this(indexConfigs, enableDefaultStarTree, indexDir, buildMode, null);
+ }
+
+ /**
+ * Constructor for star-tree build invoked during segment creation. Derives
the star-tree index
+ * configs, enable-default flag, build mode, and (optional)
IndexLoadingConfig from the supplied
+ * {@link SegmentGeneratorConfig} so downstream readers (e.g. external-table
forward-index readers
+ * backed by remote storage) can resolve table-level configs during the
in-place segment load.
+ */
+ public MultipleTreesBuilder(File indexDir, SegmentGeneratorConfig
segmentGeneratorConfig)
+ throws Exception {
+ this(segmentGeneratorConfig.getStarTreeIndexConfigs(),
segmentGeneratorConfig.isEnableDefaultStarTree(), indexDir,
+ segmentGeneratorConfig.isOnHeap() ? BuildMode.ON_HEAP :
BuildMode.OFF_HEAP,
+ (segmentGeneratorConfig.getTableConfig() != null &&
segmentGeneratorConfig.getSchema() != null)
+ ? new IndexLoadingConfig(segmentGeneratorConfig.getTableConfig(),
segmentGeneratorConfig.getSchema())
+ : null);
+ }
+
+ public MultipleTreesBuilder(@Nullable List<StarTreeIndexConfig>
indexConfigs, boolean enableDefaultStarTree,
+ File indexDir, BuildMode buildMode, @Nullable IndexLoadingConfig
indexLoadingConfig)
+ throws Exception {
Preconditions.checkArgument(CollectionUtils.isNotEmpty(indexConfigs) ||
enableDefaultStarTree,
"Must provide star-tree index configs or enable default star-tree");
_buildMode = buildMode;
@@ -126,7 +171,7 @@ public class MultipleTreesBuilder implements Closeable {
_metadataProperties =
CommonsConfigurationUtils.fromFile(new File(_segmentDirectory,
V1Constants.MetadataKeys.METADATA_FILE_NAME));
Preconditions.checkState(!_metadataProperties.containsKey(MetadataKey.STAR_TREE_COUNT),
"Star-tree already exists");
- _segment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap);
+ _segment = loadSegment(indexDir, indexLoadingConfig);
try {
_builderConfigs =
StarTreeBuilderUtils.generateBuilderConfigs(indexConfigs, enableDefaultStarTree,
_segment.getSegmentMetadata());
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java
index cf0d4175141..ded384a922f 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java
@@ -2297,8 +2297,9 @@ public class ForwardIndexHandlerTest {
if (expectDictionary && forwardIndexReader.isDictionaryEncoded()) {
dictionary = DictionaryIndexType.read(reader, columnMetadata);
}
- PinotSegmentColumnReader columnReader = new
PinotSegmentColumnReader(forwardIndexReader, dictionary, null,
- columnMetadata.getMaxNumberOfMultiValues());
+ PinotSegmentColumnReader columnReader =
+ new PinotSegmentColumnReader(columnMetadata.getColumnName(),
forwardIndexReader, dictionary, null,
+ columnMetadata.getMaxNumberOfMultiValues());
if (expectDictionary && !forwardIndexReader.isDictionaryEncoded()) {
try (Dictionary loadedDictionary = DictionaryIndexType.read(reader,
columnMetadata)) {
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderRawWithDictionaryTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderRawWithDictionaryTest.java
new file mode 100644
index 00000000000..85c25696644
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderRawWithDictionaryTest.java
@@ -0,0 +1,148 @@
+/**
+ * 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.segment.readers;
+
+import java.math.BigDecimal;
+import org.apache.pinot.segment.spi.index.reader.Dictionary;
+import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Unit tests for {@link PinotSegmentColumnReader#getDictId(int)} when the
forward index is RAW
+/// and a separate (shared) dictionary is materialized for the column. This is
the encoding used by
+/// star-tree dimension columns that opt into shared dictionaries — the FI
stores raw values, and
+/// the star-tree builder needs to resolve a dict-id per doc by looking the
raw value up in the
+/// dictionary.
+public class PinotSegmentColumnReaderRawWithDictionaryTest {
+
+ @DataProvider(name = "rawTypes")
+ public Object[][] rawTypes() {
+ return new Object[][]{
+ {DataType.INT, 42, 7},
+ {DataType.LONG, 9_000_000_000L, 11},
+ {DataType.FLOAT, 3.14f, 3},
+ {DataType.DOUBLE, 2.71828d, 5},
+ {DataType.BIG_DECIMAL, new BigDecimal("123.456"), 17},
+ {DataType.STRING, "hello-world", 23},
+ };
+ }
+
+ @Test(dataProvider = "rawTypes")
+ public void testGetDictIdResolvesFromDictionaryForRawForwardIndex(DataType
type, Object value,
+ int expectedDictId)
+ throws Exception {
+ ForwardIndexReader<?> fiReader = mock(ForwardIndexReader.class);
+ Dictionary dictionary = mock(Dictionary.class);
+ when(fiReader.isSingleValue()).thenReturn(true);
+ when(fiReader.isDictionaryEncoded()).thenReturn(false);
+ when(fiReader.getStoredType()).thenReturn(type);
+ when(dictionary.getValueType()).thenReturn(type);
+
+ switch (type) {
+ case INT:
+ when(fiReader.getInt(eq(0), any())).thenReturn((Integer) value);
+ when(dictionary.indexOf((int) (Integer)
value)).thenReturn(expectedDictId);
+ break;
+ case LONG:
+ when(fiReader.getLong(eq(0), any())).thenReturn((Long) value);
+ when(dictionary.indexOf((long) (Long)
value)).thenReturn(expectedDictId);
+ break;
+ case FLOAT:
+ when(fiReader.getFloat(eq(0), any())).thenReturn((Float) value);
+ when(dictionary.indexOf((float) (Float)
value)).thenReturn(expectedDictId);
+ break;
+ case DOUBLE:
+ when(fiReader.getDouble(eq(0), any())).thenReturn((Double) value);
+ when(dictionary.indexOf((double) (Double)
value)).thenReturn(expectedDictId);
+ break;
+ case BIG_DECIMAL:
+ when(fiReader.getBigDecimal(eq(0), any())).thenReturn((BigDecimal)
value);
+ when(dictionary.indexOf((BigDecimal)
value)).thenReturn(expectedDictId);
+ break;
+ case STRING:
+ when(fiReader.getString(eq(0), any())).thenReturn((String) value);
+ when(dictionary.indexOf((String) value)).thenReturn(expectedDictId);
+ break;
+ default:
+ throw new IllegalArgumentException("Unexpected type: " + type);
+ }
+
+ PinotSegmentColumnReader reader = new
PinotSegmentColumnReader("testColumn", fiReader, dictionary, null, 0);
+ assertEquals(reader.getDictId(0), expectedDictId);
+ }
+
+ @Test
+ public void testGetDictIdResolvesFromDictionaryForRawBytesForwardIndex()
+ throws Exception {
+ ForwardIndexReader<?> fiReader = mock(ForwardIndexReader.class);
+ Dictionary dictionary = mock(Dictionary.class);
+ byte[] raw = new byte[]{1, 2, 3, 4};
+ int expectedDictId = 13;
+
+ when(fiReader.isSingleValue()).thenReturn(true);
+ when(fiReader.isDictionaryEncoded()).thenReturn(false);
+ when(fiReader.getStoredType()).thenReturn(DataType.BYTES);
+ when(dictionary.getValueType()).thenReturn(DataType.BYTES);
+ when(fiReader.getBytes(eq(0), any())).thenReturn(raw);
+ when(dictionary.indexOf(new ByteArray(raw))).thenReturn(expectedDictId);
+
+ PinotSegmentColumnReader reader = new
PinotSegmentColumnReader("testColumn", fiReader, dictionary, null, 0);
+ assertEquals(reader.getDictId(0), expectedDictId);
+ }
+
+ @Test
+ public void testGetDictIdDelegatesToForwardIndexWhenDictionaryEncoded() {
+ ForwardIndexReader<?> fiReader = mock(ForwardIndexReader.class);
+ Dictionary dictionary = mock(Dictionary.class);
+ when(fiReader.isSingleValue()).thenReturn(true);
+ when(fiReader.isDictionaryEncoded()).thenReturn(true);
+ when(fiReader.getStoredType()).thenReturn(DataType.INT);
+ when(dictionary.getValueType()).thenReturn(DataType.INT);
+ when(fiReader.getDictId(eq(3), any())).thenReturn(99);
+
+ PinotSegmentColumnReader reader = new
PinotSegmentColumnReader("testColumn", fiReader, dictionary, null, 0);
+ assertEquals(reader.getDictId(3), 99);
+ verify(fiReader).getDictId(eq(3), any());
+ // No raw-value lookup should be attempted in the dictionary-encoded path.
+ verify(dictionary, org.mockito.Mockito.never()).indexOf(anyInt());
+ }
+
+ @Test
+ public void testGetDictIdThrowsWhenRawAndNoDictionaryMaterialized() {
+ ForwardIndexReader<?> fiReader = mock(ForwardIndexReader.class);
+ when(fiReader.isSingleValue()).thenReturn(true);
+ when(fiReader.isDictionaryEncoded()).thenReturn(false);
+ when(fiReader.getStoredType()).thenReturn(DataType.INT);
+
+ PinotSegmentColumnReader reader = new
PinotSegmentColumnReader("testColumn", fiReader, null, null, 0);
+ assertThrows(UnsupportedOperationException.class, () ->
reader.getDictId(0));
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
