sorted column change wrt to input data
+ *
+ * Initialize the driver using {@link #init(SegmentGeneratorConfig, ColumnReaderFactory)}
+ *
+ * @throws Exception if segment building fails
+ */
+ private void buildColumnar()
+ throws Exception {
+ if (!(_dataSource instanceof ColumnarSegmentCreationDataSource)) {
+ throw new IllegalStateException("buildColumnar() can only be called after initColumnar()");
+ }
+
+ ColumnarSegmentCreationDataSource columnarDataSource = (ColumnarSegmentCreationDataSource) _dataSource;
+
+ try {
+ Map columnReaders = columnarDataSource.getColumnReaders();
+
+ LOGGER.info("Starting columnar segment building for {} columns", columnReaders.size());
+
+ // Reuse existing stats collection and index creation info logic
+ LOGGER.debug("Start building StatsCollector!");
+ collectStatsAndIndexCreationInfo();
+ LOGGER.info("Finished building StatsCollector!");
+ LOGGER.info("Collected stats for {} documents", _totalDocs);
+
+ if (_totalDocs == 0) {
+ LOGGER.warn("No documents found in data source");
+ handlePostCreation();
+ return;
+ }
+
+ // Initialize the index creation using the per-column statistics information
+ _indexCreator.init(_config, _segmentIndexCreationInfo, _indexCreationInfoMap, _dataSchema, _tempIndexDir, null);
+
+ // Build the indexes column-wise (true column-major approach)
+ LOGGER.info("Start building Index using columnar approach");
+ long indexStartTime = System.nanoTime();
+
+ TreeSet columns = _dataSchema.getPhysicalColumnNames();
+ for (String columnName : columns) {
+ LOGGER.debug("Indexing column: {}", columnName);
+ ColumnReader columnReader = columnReaders.get(columnName);
+ if (columnReader == null) {
+ throw new IllegalStateException("No column reader found for column: " + columnName);
+ }
+
+ // Index each column independently using true column-major approach
+ // This is similar to how buildByColumn works but uses ColumnReader instead of IndexSegment
+ _indexCreator.indexColumn(columnName, columnReader);
+ }
+
+ _totalIndexTimeNs = System.nanoTime() - indexStartTime;
+
+ LOGGER.info("Finished indexing using columnar approach in IndexCreator!");
+ handlePostCreation();
+ } catch (Exception e) {
+ try {
+ _indexCreator.close();
+ } catch (Exception closeException) {
+ LOGGER.error("Error closing index creator", closeException);
+ // Add the close exception as suppressed to preserve both exceptions
+ e.addSuppressed(closeException);
+ }
+ throw e;
+ } finally {
+ // Always close the columnar data source to prevent resource leaks
+ try {
+ columnarDataSource.close();
+ } catch (Exception closeException) {
+ LOGGER.error("Error closing columnar data source", closeException);
+ }
+ }
+ }
}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java
index f5aa512c43..f4f3f302fb 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java
@@ -99,6 +99,40 @@ public boolean isSorted() {
*/
public abstract void collect(Object entry);
+ // Collects statistics for primitive values
+ // Default implementation boxes the value for backward compatibility.
+ public void collect(int value) {
+ collect((Object) value);
+ }
+
+ public void collect(long value) {
+ collect((Object) value);
+ }
+
+ public void collect(float value) {
+ collect((Object) value);
+ }
+
+ public void collect(double value) {
+ collect((Object) value);
+ }
+
+ public void collect(int[] values) {
+ collect((Object) values);
+ }
+
+ public void collect(long[] values) {
+ collect((Object) values);
+ }
+
+ public void collect(float[] values) {
+ collect((Object) values);
+ }
+
+ public void collect(double[] values) {
+ collect((Object) values);
+ }
+
public int getLengthOfShortestElement() {
return -1;
}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnarSegmentPreIndexStatsContainer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnarSegmentPreIndexStatsContainer.java
new file mode 100644
index 0000000000..88d829ae06
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnarSegmentPreIndexStatsContainer.java
@@ -0,0 +1,178 @@
+/**
+ * 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.creator.impl.stats;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.segment.spi.creator.ColumnStatistics;
+import org.apache.pinot.segment.spi.creator.SegmentPreIndexStatsCollector;
+import org.apache.pinot.segment.spi.creator.StatsCollectorConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Stats container that efficiently collects statistics from columnar data using ColumnReader instances.
+ *
+ * This implementation collects statistics by iterating column-wise instead of row-wise,
+ * which is more efficient for columnar data sources. It supports:
+ *
+ * - Column-wise statistics collection
+ * - Existing columns from source data
+ * - New columns with default values
+ * - Data type conversions during schema evolution
+ *
+ *
+ * The statistics are collected using the same underlying collectors as the row-based approach
+ * (SegmentPreIndexStatsCollectorImpl) but with more efficient column-wise iteration.
+ */
+public class ColumnarSegmentPreIndexStatsContainer implements SegmentPreIndexStatsCollector {
+ private static final Logger LOGGER = LoggerFactory.getLogger(ColumnarSegmentPreIndexStatsContainer.class);
+
+ private final Map _columnReaders;
+ private final StatsCollectorConfig _statsCollectorConfig;
+ private final Schema _targetSchema;
+ private final Map _columnStatsCollectorMap;
+ private int _totalDocCount;
+
+ /**
+ * Create a ColumnarSegmentPreIndexStatsContainer.
+ *
+ * @param columnReaders Map of column name to ColumnReader instances
+ * @param statsCollectorConfig Configuration for statistics collection
+ */
+ public ColumnarSegmentPreIndexStatsContainer(Map columnReaders,
+ StatsCollectorConfig statsCollectorConfig) {
+ _columnReaders = columnReaders;
+ _statsCollectorConfig = statsCollectorConfig;
+ _targetSchema = statsCollectorConfig.getSchema();
+ _columnStatsCollectorMap = new HashMap<>(columnReaders.size());
+ _totalDocCount = -1; // indicates unset
+ }
+
+ /**
+ * Initialize stats collectors for all columns in the target schema.
+ */
+ private void initializeStatsCollectors() {
+ for (FieldSpec fieldSpec : _targetSchema.getAllFieldSpecs()) {
+ if (fieldSpec.isVirtualColumn()) {
+ continue;
+ }
+
+ String columnName = fieldSpec.getName();
+ AbstractColumnStatisticsCollector collector =
+ StatsCollectorUtil.createStatsCollector(columnName, fieldSpec, _statsCollectorConfig);
+ _columnStatsCollectorMap.put(columnName, collector);
+ }
+ }
+
+ /**
+ * Collect stats by iterating column-wise using the provided ColumnReader instances.
+ */
+ private void collectColumnStats() {
+ LOGGER.info("Collecting stats for {} columns using column-wise iteration", _columnReaders.size());
+
+ for (String columnName : _columnStatsCollectorMap.keySet()) {
+ AbstractColumnStatisticsCollector statsCollector = _columnStatsCollectorMap.get(columnName);
+ ColumnReader columnReader = _columnReaders.get(columnName);
+
+ if (columnReader == null) {
+ throw new RuntimeException("Column reader for column " + columnName + " not found");
+ }
+
+ LOGGER.debug("Collecting stats for column: {}", columnName);
+ collectStatsFromColumnReader(columnName, columnReader, statsCollector);
+
+ // Seal the stats collector
+ statsCollector.seal();
+ }
+ }
+
+ /**
+ * Collect stats from a column reader by iterating over all values using the iterator pattern.
+ */
+ private void collectStatsFromColumnReader(String columnName, ColumnReader columnReader,
+ AbstractColumnStatisticsCollector statsCollector) {
+ try {
+ // Reset the column reader to start from the beginning
+ columnReader.rewind();
+
+ int docCount = 0;
+ while (columnReader.hasNext()) {
+ Object value = columnReader.next();
+ statsCollector.collect(value);
+ docCount++;
+ }
+
+ // if totalDocCount is unset then set total doc count from the first column
+ if (_totalDocCount == -1) {
+ _totalDocCount = docCount;
+ } else if (_totalDocCount != docCount) {
+ // all columns should have same count
+ LOGGER.warn("Column {} has {} documents, but expected {} documents",
+ columnName, docCount, _totalDocCount);
+ throw new RuntimeException("Columns have inconsistent document counts");
+ }
+ } catch (IOException e) {
+ LOGGER.error("Failed to collect stats for column: {}", columnName, e);
+ throw new RuntimeException("Failed to collect stats for column: " + columnName, e);
+ }
+ }
+
+ @Override
+ public ColumnStatistics getColumnProfileFor(String column) {
+ return _columnStatsCollectorMap.get(column);
+ }
+
+ @Override
+ public int getTotalDocCount() {
+ return _totalDocCount;
+ }
+
+ @Override
+ public void init() {
+ initializeStatsCollectors();
+ collectColumnStats();
+ }
+
+ @Override
+ public void build() {
+ // Stats are already collected in constructor
+ }
+
+ @Override
+ public void logStats() {
+ LOGGER.info("Columnar segment stats collection completed for {} columns with {} documents",
+ _columnStatsCollectorMap.size(), _totalDocCount);
+ }
+
+ @Override
+ public void collectRow(GenericRow row) {
+ // This method is not used in columnar stats collection as we iterate column-wise
+ // instead of row-wise. The stats are collected in the constructor by iterating
+ // over each column using ColumnReader instances.
+ throw new UnsupportedOperationException(
+ "collectRow is not supported in columnar stats collection. Stats are collected column-wise.");
+ }
+}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/DoubleColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/DoubleColumnPreIndexStatsCollector.java
index 9edd89b027..77b34b0459 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/DoubleColumnPreIndexStatsCollector.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/DoubleColumnPreIndexStatsCollector.java
@@ -40,32 +40,38 @@ public void collect(Object entry) {
if (entry instanceof Object[]) {
Object[] values = (Object[]) entry;
- for (Object obj : values) {
- double value = (double) obj;
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values);
+ collectMultiValue(values.length, i -> (double) values[i]);
} else if (entry instanceof double[]) {
- double[] values = (double[]) entry;
- for (double value : values) {
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values.length);
+ collect((double[]) entry);
} else {
- double value = (double) entry;
- addressSorted(value);
- if (_values.add(value)) {
- if (isPartitionEnabled()) {
- updatePartition(Double.toString(value));
- }
+ collect((double) entry);
+ }
+ }
+
+ @Override
+ public void collect(double value) {
+ assert !_sealed;
+ addressSorted(value);
+ if (_values.add(value)) {
+ if (isPartitionEnabled()) {
+ updatePartition(Double.toString(value));
}
+ }
+ _totalNumberOfEntries++;
+ }
+
+ @Override
+ public void collect(double[] values) {
+ assert !_sealed;
+ collectMultiValue(values.length, i -> values[i]);
+ }
- _totalNumberOfEntries++;
+ private void collectMultiValue(int length, java.util.function.IntFunction valueGetter) {
+ for (int i = 0; i < length; i++) {
+ _values.add(valueGetter.apply(i));
}
+ _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, length);
+ updateTotalNumberOfEntries(length);
}
private void addressSorted(double entry) {
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/FloatColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/FloatColumnPreIndexStatsCollector.java
index 4be248124e..aa60f4d5ba 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/FloatColumnPreIndexStatsCollector.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/FloatColumnPreIndexStatsCollector.java
@@ -40,32 +40,38 @@ public void collect(Object entry) {
if (entry instanceof Object[]) {
Object[] values = (Object[]) entry;
- for (Object obj : values) {
- float value = (float) obj;
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values);
+ collectMultiValue(values.length, i -> (float) values[i]);
} else if (entry instanceof float[]) {
- float[] values = (float[]) entry;
- for (float value : values) {
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values.length);
+ collect((float[]) entry);
} else {
- float value = (float) entry;
- addressSorted(value);
- if (_values.add(value)) {
- if (isPartitionEnabled()) {
- updatePartition(Float.toString(value));
- }
+ collect((float) entry);
+ }
+ }
+
+ @Override
+ public void collect(float value) {
+ assert !_sealed;
+ addressSorted(value);
+ if (_values.add(value)) {
+ if (isPartitionEnabled()) {
+ updatePartition(Float.toString(value));
}
+ }
+ _totalNumberOfEntries++;
+ }
+
+ @Override
+ public void collect(float[] values) {
+ assert !_sealed;
+ collectMultiValue(values.length, i -> values[i]);
+ }
- _totalNumberOfEntries++;
+ private void collectMultiValue(int length, java.util.function.IntFunction valueGetter) {
+ for (int i = 0; i < length; i++) {
+ _values.add(valueGetter.apply(i));
}
+ _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, length);
+ updateTotalNumberOfEntries(length);
}
private void addressSorted(float entry) {
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/IntColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/IntColumnPreIndexStatsCollector.java
index 29678f6805..74882f25ff 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/IntColumnPreIndexStatsCollector.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/IntColumnPreIndexStatsCollector.java
@@ -40,32 +40,38 @@ public void collect(Object entry) {
if (entry instanceof Object[]) {
Object[] values = (Object[]) entry;
- for (Object obj : values) {
- int value = (int) obj;
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values);
+ collectMultiValue(values.length, i -> (int) values[i]);
} else if (entry instanceof int[]) {
- int[] values = (int[]) entry;
- for (int value : values) {
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values.length);
+ collect((int[]) entry);
} else {
- int value = (int) entry;
- addressSorted(value);
- if (_values.add(value)) {
- if (isPartitionEnabled()) {
- updatePartition(Integer.toString(value));
- }
+ collect((int) entry);
+ }
+ }
+
+ @Override
+ public void collect(int value) {
+ assert !_sealed;
+ addressSorted(value);
+ if (_values.add(value)) {
+ if (isPartitionEnabled()) {
+ updatePartition(Integer.toString(value));
}
+ }
+ _totalNumberOfEntries++;
+ }
+
+ @Override
+ public void collect(int[] values) {
+ assert !_sealed;
+ collectMultiValue(values.length, i -> values[i]);
+ }
- _totalNumberOfEntries++;
+ private void collectMultiValue(int length, java.util.function.IntFunction valueGetter) {
+ for (int i = 0; i < length; i++) {
+ _values.add(valueGetter.apply(i));
}
+ _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, length);
+ updateTotalNumberOfEntries(length);
}
private void addressSorted(int entry) {
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/LongColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/LongColumnPreIndexStatsCollector.java
index 3d4c63d03b..c60efc5bd9 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/LongColumnPreIndexStatsCollector.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/LongColumnPreIndexStatsCollector.java
@@ -40,32 +40,38 @@ public void collect(Object entry) {
if (entry instanceof Object[]) {
Object[] values = (Object[]) entry;
- for (Object obj : values) {
- long value = (long) obj;
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values);
+ collectMultiValue(values.length, i -> (long) values[i]);
} else if (entry instanceof long[]) {
- long[] values = (long[]) entry;
- for (long value : values) {
- _values.add(value);
- }
-
- _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, values.length);
- updateTotalNumberOfEntries(values.length);
+ collect((long[]) entry);
} else {
- long value = (long) entry;
- addressSorted(value);
- if (_values.add(value)) {
- if (isPartitionEnabled()) {
- updatePartition(Long.toString(value));
- }
+ collect((long) entry);
+ }
+ }
+
+ @Override
+ public void collect(long value) {
+ assert !_sealed;
+ addressSorted(value);
+ if (_values.add(value)) {
+ if (isPartitionEnabled()) {
+ updatePartition(Long.toString(value));
}
+ }
+ _totalNumberOfEntries++;
+ }
+
+ @Override
+ public void collect(long[] values) {
+ assert !_sealed;
+ collectMultiValue(values.length, i -> values[i]);
+ }
- _totalNumberOfEntries++;
+ private void collectMultiValue(int length, java.util.function.IntFunction valueGetter) {
+ for (int i = 0; i < length; i++) {
+ _values.add(valueGetter.apply(i));
}
+ _maxNumberOfMultiValues = Math.max(_maxNumberOfMultiValues, length);
+ updateTotalNumberOfEntries(length);
}
private void addressSorted(long entry) {
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java
index 45c75e0c49..0a680550a9 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java
@@ -48,35 +48,9 @@ public void init() {
Schema dataSchema = _statsCollectorConfig.getSchema();
for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) {
String column = fieldSpec.getName();
- switch (fieldSpec.getDataType().getStoredType()) {
- case INT:
- _columnStatsCollectorMap.put(column, new IntColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case LONG:
- _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case FLOAT:
- _columnStatsCollectorMap.put(column, new FloatColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case DOUBLE:
- _columnStatsCollectorMap.put(column, new DoubleColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case BIG_DECIMAL:
- _columnStatsCollectorMap.put(column,
- new BigDecimalColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case STRING:
- _columnStatsCollectorMap.put(column, new StringColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case BYTES:
- _columnStatsCollectorMap.put(column, new BytesColumnPredIndexStatsCollector(column, _statsCollectorConfig));
- break;
- case MAP:
- _columnStatsCollectorMap.put(column, new MapColumnPreIndexStatsCollector(column, _statsCollectorConfig));
- break;
- default:
- throw new IllegalStateException("Unsupported data type: " + fieldSpec.getDataType());
- }
+ AbstractColumnStatisticsCollector collector =
+ StatsCollectorUtil.createStatsCollector(column, fieldSpec, _statsCollectorConfig);
+ _columnStatsCollectorMap.put(column, collector);
}
}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StatsCollectorUtil.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StatsCollectorUtil.java
new file mode 100644
index 0000000000..469184c3f4
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StatsCollectorUtil.java
@@ -0,0 +1,65 @@
+/**
+ * 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.creator.impl.stats;
+
+import org.apache.pinot.segment.spi.creator.StatsCollectorConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+
+
+/**
+ * Utility class for creating column statistics collectors.
+ */
+public final class StatsCollectorUtil {
+
+ private StatsCollectorUtil() {
+ // Utility class
+ }
+
+ /**
+ * Create a statistics collector for the given column based on its data type.
+ *
+ * @param columnName Name of the column
+ * @param fieldSpec Field specification for the column
+ * @param statsCollectorConfig Stats collector configuration
+ * @return AbstractColumnStatisticsCollector for the column
+ */
+ public static AbstractColumnStatisticsCollector createStatsCollector(String columnName, FieldSpec fieldSpec,
+ StatsCollectorConfig statsCollectorConfig) {
+ switch (fieldSpec.getDataType().getStoredType()) {
+ case INT:
+ return new IntColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case LONG:
+ return new LongColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case FLOAT:
+ return new FloatColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case DOUBLE:
+ return new DoubleColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case BIG_DECIMAL:
+ return new BigDecimalColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case STRING:
+ return new StringColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ case BYTES:
+ return new BytesColumnPredIndexStatsCollector(columnName, statsCollectorConfig);
+ case MAP:
+ return new MapColumnPreIndexStatsCollector(columnName, statsCollectorConfig);
+ default:
+ throw new IllegalStateException("Unsupported data type: " + fieldSpec.getDataType());
+ }
+ }
+}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java
new file mode 100644
index 0000000000..4392fe9cd0
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java
@@ -0,0 +1,410 @@
+/**
+ * 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 javax.annotation.Nullable;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.data.readers.MultiValueResult;
+
+
+/**
+ * ColumnReader implementation that returns default values for new columns.
+ *
+ * This reader is used when a column exists in the target schema but not in the source data.
+ * It returns the default null value for the field spec for all document IDs.
+ */
+public class DefaultValueColumnReader implements ColumnReader {
+
+ private final String _columnName;
+ private final int _numDocs;
+ private final Object _defaultValue;
+ private final FieldSpec _fieldSpec;
+ private final FieldSpec.DataType _dataType;
+
+ // Pre-computed multi-value arrays for reuse
+ private int[] _defaultIntMV;
+ private long[] _defaultLongMV;
+ private float[] _defaultFloatMV;
+ private double[] _defaultDoubleMV;
+ private String[] _defaultStringMV;
+ private byte[][] _defaultBytesMV;
+
+ private int _currentIndex;
+
+ /**
+ * Create a DefaultValueColumnReader for a new column.
+ *
+ * @param columnName Name of the new column
+ * @param numDocs Total number of documents
+ * @param fieldSpec Field specification for the new column
+ */
+ public DefaultValueColumnReader(String columnName, int numDocs, FieldSpec fieldSpec) {
+ _columnName = columnName;
+ _numDocs = numDocs;
+ _currentIndex = 0;
+ _fieldSpec = fieldSpec;
+ _dataType = fieldSpec.getDataType();
+
+ // For multi-value fields, wrap the default value in an array
+ Object defaultNullValue = fieldSpec.getDefaultNullValue();
+ if (fieldSpec.isSingleValueField()) {
+ _defaultValue = defaultNullValue;
+ } else {
+ _defaultValue = new Object[]{defaultNullValue};
+ // Pre-compute typed arrays for multi-value fields to avoid repeated allocations
+ Object[] defaultArray = (Object[]) _defaultValue;
+ switch (_dataType) {
+ case INT:
+ _defaultIntMV = new int[defaultArray.length];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultIntMV[i] = ((Number) defaultArray[i]).intValue();
+ }
+ break;
+ case LONG:
+ _defaultLongMV = new long[defaultArray.length];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultLongMV[i] = ((Number) defaultArray[i]).longValue();
+ }
+ break;
+ case FLOAT:
+ _defaultFloatMV = new float[defaultArray.length];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultFloatMV[i] = ((Number) defaultArray[i]).floatValue();
+ }
+ break;
+ case DOUBLE:
+ _defaultDoubleMV = new double[defaultArray.length];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultDoubleMV[i] = ((Number) defaultArray[i]).doubleValue();
+ }
+ break;
+ case STRING:
+ _defaultStringMV = new String[defaultArray.length];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultStringMV[i] = (String) defaultArray[i];
+ }
+ break;
+ case BYTES:
+ _defaultBytesMV = new byte[defaultArray.length][];
+ for (int i = 0; i < defaultArray.length; i++) {
+ _defaultBytesMV[i] = (byte[]) defaultArray[i];
+ }
+ break;
+ default:
+ break;
+ }
+ }
+ }
+
+ @Override
+ public boolean hasNext() {
+ return _currentIndex < _numDocs;
+ }
+
+ @Override
+ @Nullable
+ public Object next() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return _defaultValue;
+ }
+
+ @Override
+ public void rewind() {
+ _currentIndex = 0;
+ }
+
+ @Override
+ public boolean isNextNull() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ return _defaultValue == null;
+ }
+
+ @Override
+ public void skipNext() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ }
+
+ @Override
+ public boolean isSingleValue() {
+ return _fieldSpec.isSingleValueField();
+ }
+
+ @Override
+ public boolean isInt() {
+ return _dataType == FieldSpec.DataType.INT;
+ }
+
+ @Override
+ public boolean isLong() {
+ return _dataType == FieldSpec.DataType.LONG;
+ }
+
+ @Override
+ public boolean isFloat() {
+ return _dataType == FieldSpec.DataType.FLOAT;
+ }
+
+ @Override
+ public boolean isDouble() {
+ return _dataType == FieldSpec.DataType.DOUBLE;
+ }
+
+ @Override
+ public boolean isString() {
+ return _dataType == FieldSpec.DataType.STRING;
+ }
+
+ @Override
+ public boolean isBytes() {
+ return _dataType == FieldSpec.DataType.BYTES;
+ }
+
+ @Override
+ public int nextInt() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return ((Number) _defaultValue).intValue();
+ }
+
+ @Override
+ public long nextLong() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return ((Number) _defaultValue).longValue();
+ }
+
+ @Override
+ public float nextFloat() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return ((Number) _defaultValue).floatValue();
+ }
+
+ @Override
+ public double nextDouble() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return ((Number) _defaultValue).doubleValue();
+ }
+
+ @Override
+ public String nextString() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return (String) _defaultValue;
+ }
+
+ @Override
+ public byte[] nextBytes() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return (byte[]) _defaultValue;
+ }
+
+ @Override
+ public MultiValueResult nextIntMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return MultiValueResult.of(_defaultIntMV, null);
+ }
+
+ @Override
+ public MultiValueResult nextLongMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return MultiValueResult.of(_defaultLongMV, null);
+ }
+
+ @Override
+ public MultiValueResult nextFloatMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return MultiValueResult.of(_defaultFloatMV, null);
+ }
+
+ @Override
+ public MultiValueResult nextDoubleMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return MultiValueResult.of(_defaultDoubleMV, null);
+ }
+
+ @Override
+ public String[] nextStringMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return _defaultStringMV;
+ }
+
+ @Override
+ public byte[][] nextBytesMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ return _defaultBytesMV;
+ }
+
+ @Override
+ public String getColumnName() {
+ return _columnName;
+ }
+
+ @Override
+ public int getTotalDocs() {
+ return _numDocs;
+ }
+
+ @Override
+ public boolean isNull(int docId) {
+ validateDocId(docId);
+ return _defaultValue == null;
+ }
+
+ // Single-value accessors
+
+ @Override
+ public int getInt(int docId) {
+ validateDocId(docId);
+ return ((Number) _defaultValue).intValue();
+ }
+
+ @Override
+ public long getLong(int docId) {
+ validateDocId(docId);
+ return ((Number) _defaultValue).longValue();
+ }
+
+ @Override
+ public float getFloat(int docId) {
+ validateDocId(docId);
+ return ((Number) _defaultValue).floatValue();
+ }
+
+ @Override
+ public double getDouble(int docId) {
+ validateDocId(docId);
+ return ((Number) _defaultValue).doubleValue();
+ }
+
+ @Override
+ public String getString(int docId) {
+ validateDocId(docId);
+ return (String) _defaultValue;
+ }
+
+ @Override
+ public byte[] getBytes(int docId) {
+ validateDocId(docId);
+ return (byte[]) _defaultValue;
+ }
+
+ @Override
+ public Object getValue(int docId) {
+ validateDocId(docId);
+ return _defaultValue;
+ }
+
+ // Multi-value accessors
+
+ @Override
+ public MultiValueResult getIntMV(int docId) {
+ validateDocId(docId);
+ return MultiValueResult.of(_defaultIntMV, null);
+ }
+
+ @Override
+ public MultiValueResult getLongMV(int docId) {
+ validateDocId(docId);
+ return MultiValueResult.of(_defaultLongMV, null);
+ }
+
+ @Override
+ public MultiValueResult getFloatMV(int docId) {
+ validateDocId(docId);
+ return MultiValueResult.of(_defaultFloatMV, null);
+ }
+
+ @Override
+ public MultiValueResult getDoubleMV(int docId) {
+ validateDocId(docId);
+ return MultiValueResult.of(_defaultDoubleMV, null);
+ }
+
+ @Override
+ public String[] getStringMV(int docId) {
+ validateDocId(docId);
+ return _defaultStringMV;
+ }
+
+ @Override
+ public byte[][] getBytesMV(int docId) {
+ validateDocId(docId);
+ return _defaultBytesMV;
+ }
+
+ /**
+ * Validate that the document ID is within valid range.
+ *
+ * @param docId Document ID to validate
+ * @throws IndexOutOfBoundsException if docId is out of range
+ */
+ private void validateDocId(int docId) {
+ if (docId < 0 || docId >= _numDocs) {
+ throw new IndexOutOfBoundsException(
+ "docId " + docId + " is out of range. Valid range is 0 to " + (_numDocs - 1));
+ }
+ }
+
+ @Override
+ public void close() {
+ // No resources to close
+ }
+}
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 4a691e4fc4..4e62b6e5af 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
@@ -177,6 +177,124 @@ public boolean isNull(int docId) {
return _nullValueVectorReader != null && _nullValueVectorReader.isNull(docId);
}
+ // Single-value accessors
+
+ public int getInt(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getIntValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getInt(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public long getLong(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getLongValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getLong(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public float getFloat(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getFloatValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getFloat(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public double getDouble(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getDoubleValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getDouble(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public String getString(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getStringValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getString(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public byte[] getBytes(int docId) {
+ if (_dictionary != null) {
+ return _dictionary.getBytesValue(_forwardIndexReader.getDictId(docId, _forwardIndexReaderContext));
+ } else {
+ return _forwardIndexReader.getBytes(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ // Multi-value accessors
+
+ public int[] getIntMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ int[] values = new int[numValues];
+ _dictionary.readIntValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getIntMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public long[] getLongMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ long[] values = new long[numValues];
+ _dictionary.readLongValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getLongMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public float[] getFloatMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ float[] values = new float[numValues];
+ _dictionary.readFloatValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getFloatMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public double[] getDoubleMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ double[] values = new double[numValues];
+ _dictionary.readDoubleValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getDoubleMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public String[] getStringMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ String[] values = new String[numValues];
+ _dictionary.readStringValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getStringMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
+ public byte[][] getBytesMV(int docId) {
+ if (_dictionary != null) {
+ int numValues = _forwardIndexReader.getDictIdMV(docId, _dictIdBuffer, _forwardIndexReaderContext);
+ byte[][] values = new byte[numValues][];
+ _dictionary.readBytesValues(_dictIdBuffer, numValues, values);
+ return values;
+ } else {
+ return _forwardIndexReader.getBytesMV(docId, _forwardIndexReaderContext);
+ }
+ }
+
@Override
public void close()
throws IOException {
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderFactory.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderFactory.java
new file mode 100644
index 0000000000..b6eb206c31
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderFactory.java
@@ -0,0 +1,202 @@
+/**
+ * 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.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.data.readers.ColumnReaderFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * ColumnReaderFactory implementation for immutable Pinot segments.
+ *
+ * This factory creates ColumnReader instances for reading data from Pinot segments
+ * in a columnar fashion. It handles:
+ *
+ * - Creating readers for existing columns in the segment
+ * - Creating default value readers for new columns
+ * - Resource management for all created readers
+ *
+ */
+public class PinotSegmentColumnReaderFactory implements ColumnReaderFactory {
+ private static final Logger LOGGER = LoggerFactory.getLogger(PinotSegmentColumnReaderFactory.class);
+
+ private final IndexSegment _indexSegment;
+ private Schema _targetSchema;
+ private Set _colsToRead;
+ private Map _columnReaders;
+ private final boolean _initializeDefaultValueReaders;
+
+ /**
+ * Create a PinotSegmentColumnReaderFactory.
+ *
+ * @param indexSegment Source segment to read from
+ */
+ public PinotSegmentColumnReaderFactory(IndexSegment indexSegment) {
+ this(indexSegment, true);
+ }
+
+ /**
+ * Create a PinotSegmentColumnReaderFactory.
+ *
+ * @param indexSegment Source segment to read from
+ * @param initializeDefaultValueReaders Whether to initialize default value readers for missing columns
+ * TODO - Ideally this factory shouldn't initialize default value readers.
+ * The clients of this factory should decide whether to create default value readers or not.
+ * This parameter is kept for backward compatibility and will be removed in future.
+ */
+ public PinotSegmentColumnReaderFactory(IndexSegment indexSegment, boolean initializeDefaultValueReaders) {
+ _indexSegment = indexSegment;
+ _columnReaders = new HashMap<>();
+ _initializeDefaultValueReaders = initializeDefaultValueReaders;
+ }
+
+ @Override
+ public void init(Schema targetSchema)
+ throws IOException {
+ init(targetSchema, targetSchema.getPhysicalColumnNames());
+ }
+
+ @Override
+ public void init(Schema targetSchema, Set colsToRead)
+ throws IOException {
+ _targetSchema = targetSchema;
+ _colsToRead = colsToRead;
+ _columnReaders = initializeAllColumnReaders();
+ LOGGER.info("Initialized PinotSegmentColumnReaderFactory with target schema containing {} columns",
+ targetSchema.getPhysicalColumnNames().size());
+ }
+
+ @Override
+ public Set getAvailableColumns() {
+ return _indexSegment.getPhysicalColumnNames();
+ }
+
+ /**
+ * Get the total number of documents/rows in the data source.
+ *
+ * @return Total number of documents
+ */
+ public int getNumDocs() {
+ return _indexSegment.getSegmentMetadata().getTotalDocs();
+ }
+
+ @Override
+ public ColumnReader getColumnReader(String columnName) {
+ if (_targetSchema == null) {
+ throw new IllegalStateException("Factory not initialized. Call init() first.");
+ }
+
+ return _columnReaders.get(columnName);
+ }
+
+ /**
+ * Internal method to create a column reader for the specified column.
+ * This method is called during initialization to create all readers.
+ */
+ @Nullable
+ private ColumnReader createColumnReader(String columnName, FieldSpec targetFieldSpec) {
+ if (targetFieldSpec.isVirtualColumn()) {
+ throw new IllegalStateException("Target field spec is a virtual column.");
+ }
+
+ ColumnReader columnReader = null;
+
+ if (hasColumn(columnName)) {
+ // Column exists in source segment - create a segment column reader
+ LOGGER.debug("Creating segment column reader for existing column: {}", columnName);
+ columnReader = new PinotSegmentColumnReaderImpl(_indexSegment, columnName);
+ } else {
+ if (_initializeDefaultValueReaders) {
+ // New column - create a default value reader
+ LOGGER.debug("Creating default value reader for new column: {}", columnName);
+ columnReader = new DefaultValueColumnReader(columnName, getNumDocs(), targetFieldSpec);
+ }
+ }
+
+ return columnReader;
+ }
+
+ @Override
+ public Map getAllColumnReaders()
+ throws IOException {
+ return _columnReaders;
+ }
+
+ /**
+ * Internal method to initialize all column readers during factory initialization.
+ */
+ private Map initializeAllColumnReaders() {
+ if (_targetSchema == null) {
+ throw new IllegalStateException("Factory not initialized. Call init() first.");
+ }
+
+ Map allReaders = new HashMap<>();
+
+ // Create readers for all columns in the target schema
+ for (String columnName : _colsToRead) {
+ FieldSpec fieldSpec = _targetSchema.getFieldSpecFor(columnName);
+ if (fieldSpec.isVirtualColumn()) {
+ continue;
+ }
+ ColumnReader reader = createColumnReader(columnName, fieldSpec);
+ if (reader != null) {
+ allReaders.put(columnName, reader);
+ }
+ }
+
+ return allReaders;
+ }
+
+ /**
+ * Check if the specified column exists in the source data.
+ *
+ * @param columnName Column name to check
+ * @return true if the column exists in source, false otherwise
+ */
+ public boolean hasColumn(String columnName) {
+ return _indexSegment.getPhysicalColumnNames().contains(columnName);
+ }
+
+ @Override
+ public void close()
+ throws IOException {
+ LOGGER.debug("Closing PinotSegmentColumnReaderFactory and {} column readers", _columnReaders.size());
+
+ // Close all created column readers
+ for (ColumnReader reader : _columnReaders.values()) {
+ try {
+ reader.close();
+ } catch (IOException e) {
+ LOGGER.warn("Error closing column reader for column: {}", reader.getColumnName(), e);
+ }
+ }
+
+ _columnReaders.clear();
+ }
+}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java
new file mode 100644
index 0000000000..b34a0e53ef
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java
@@ -0,0 +1,382 @@
+/**
+ * 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.io.IOException;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.data.readers.MultiValueResult;
+
+
+/**
+ * Implementation of ColumnReader for Pinot segments.
+ *
+ * This class wraps the existing PinotSegmentColumnReader and provides the ColumnReader interface
+ * for columnar segment building. It handles:
+ *
+ * - Reading column values from Pinot segments
+ * - Resource cleanup
+ *
+ */
+public class PinotSegmentColumnReaderImpl implements ColumnReader {
+ private final PinotSegmentColumnReader _segmentColumnReader;
+ private final String _columnName;
+ private final int _numDocs;
+ private final FieldSpec.DataType _dataType;
+
+ private int _currentIndex;
+
+ // Reusable variables to avoid garbage collection on every next() call
+ private Object _reuseValue;
+
+ /**
+ * Create a PinotSegmentColumnReaderImpl for an existing column in the segment.
+ *
+ * @param indexSegment Source segment to read from
+ * @param columnName Name of the column
+ */
+ public PinotSegmentColumnReaderImpl(IndexSegment indexSegment, String columnName) {
+ _segmentColumnReader = new PinotSegmentColumnReader(indexSegment, columnName);
+ _columnName = columnName;
+ _numDocs = indexSegment.getSegmentMetadata().getTotalDocs();
+ _currentIndex = 0;
+
+ // Get the data type from the schema
+ FieldSpec fieldSpec = indexSegment.getSegmentMetadata().getSchema().getFieldSpecFor(columnName);
+ assert fieldSpec != null : "FieldSpec should not be null for column: " + columnName;
+ _dataType = fieldSpec.getDataType();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return _currentIndex < _numDocs;
+ }
+
+ @Override
+ @Nullable
+ public Object next()
+ throws IOException {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+
+ // Return null if the value is null
+ if (_segmentColumnReader.isNull(_currentIndex)) {
+ _currentIndex++;
+ return null;
+ }
+ _reuseValue = _segmentColumnReader.getValue(_currentIndex);
+ _currentIndex++;
+
+ // Return null if the value is null
+ if (_reuseValue == null) {
+ return null;
+ }
+
+ return _reuseValue;
+ }
+
+ @Override
+ public boolean isNextNull() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ return _segmentColumnReader.isNull(_currentIndex);
+ }
+
+ @Override
+ public void skipNext() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ _currentIndex++;
+ }
+
+ @Override
+ public boolean isSingleValue() {
+ return _segmentColumnReader.isSingleValue();
+ }
+
+ @Override
+ public boolean isInt() {
+ return _dataType == FieldSpec.DataType.INT;
+ }
+
+ @Override
+ public boolean isLong() {
+ return _dataType == FieldSpec.DataType.LONG;
+ }
+
+ @Override
+ public boolean isFloat() {
+ return _dataType == FieldSpec.DataType.FLOAT;
+ }
+
+ @Override
+ public boolean isDouble() {
+ return _dataType == FieldSpec.DataType.DOUBLE;
+ }
+
+ @Override
+ public boolean isString() {
+ return _dataType == FieldSpec.DataType.STRING;
+ }
+
+ @Override
+ public boolean isBytes() {
+ return _dataType == FieldSpec.DataType.BYTES;
+ }
+
+ @Override
+ public int nextInt() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ int value = _segmentColumnReader.getInt(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public long nextLong() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ long value = _segmentColumnReader.getLong(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public float nextFloat() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ float value = _segmentColumnReader.getFloat(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public double nextDouble() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ double value = _segmentColumnReader.getDouble(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public String nextString() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ String value = _segmentColumnReader.getString(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public byte[] nextBytes() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ byte[] value = _segmentColumnReader.getBytes(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ // For all multi-value primitive type methods (nextIntMV, nextLongMV, nextFloatMV, nextDoubleMV,
+ // getIntMV, getLongMV, getFloatMV, getDoubleMV), we pass null for the validity bitset since
+ // multi-value primitive types cannot have null elements. Nulls are removed by NullValueTransformer
+ @Override
+ public MultiValueResult nextIntMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ int[] value = _segmentColumnReader.getIntMV(_currentIndex);
+ _currentIndex++;
+ return MultiValueResult.of(value, null);
+ }
+
+ @Override
+ public MultiValueResult nextLongMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ long[] value = _segmentColumnReader.getLongMV(_currentIndex);
+ _currentIndex++;
+ return MultiValueResult.of(value, null);
+ }
+
+ @Override
+ public MultiValueResult nextFloatMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ float[] value = _segmentColumnReader.getFloatMV(_currentIndex);
+ _currentIndex++;
+ return MultiValueResult.of(value, null);
+ }
+
+ @Override
+ public MultiValueResult nextDoubleMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ double[] value = _segmentColumnReader.getDoubleMV(_currentIndex);
+ _currentIndex++;
+ return MultiValueResult.of(value, null);
+ }
+
+ @Override
+ public String[] nextStringMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ String[] value = _segmentColumnReader.getStringMV(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public byte[][] nextBytesMV() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values available");
+ }
+ byte[][] value = _segmentColumnReader.getBytesMV(_currentIndex);
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public void rewind()
+ throws IOException {
+ _currentIndex = 0;
+ _reuseValue = null;
+ }
+
+ @Override
+ public String getColumnName() {
+ return _columnName;
+ }
+
+ @Override
+ public int getTotalDocs() {
+ return _numDocs;
+ }
+
+ @Override
+ public boolean isNull(int docId) {
+ validateDocId(docId);
+ return _segmentColumnReader.isNull(docId);
+ }
+
+ // Single-value accessors
+
+ @Override
+ public int getInt(int docId) {
+ return _segmentColumnReader.getInt(docId);
+ }
+
+ @Override
+ public long getLong(int docId) {
+ return _segmentColumnReader.getLong(docId);
+ }
+
+ @Override
+ public float getFloat(int docId) {
+ return _segmentColumnReader.getFloat(docId);
+ }
+
+ @Override
+ public double getDouble(int docId) {
+ return _segmentColumnReader.getDouble(docId);
+ }
+
+ @Override
+ public String getString(int docId) {
+ return _segmentColumnReader.getString(docId);
+ }
+
+ @Override
+ public byte[] getBytes(int docId) {
+ return _segmentColumnReader.getBytes(docId);
+ }
+
+ @Override
+ public Object getValue(int docId)
+ throws IOException {
+ return _segmentColumnReader.getValue(docId);
+ }
+
+ // Multi-value accessors
+
+ @Override
+ public MultiValueResult getIntMV(int docId) {
+ return MultiValueResult.of(_segmentColumnReader.getIntMV(docId), null);
+ }
+
+ @Override
+ public MultiValueResult getLongMV(int docId) {
+ return MultiValueResult.of(_segmentColumnReader.getLongMV(docId), null);
+ }
+
+ @Override
+ public MultiValueResult getFloatMV(int docId) {
+ return MultiValueResult.of(_segmentColumnReader.getFloatMV(docId), null);
+ }
+
+ @Override
+ public MultiValueResult getDoubleMV(int docId) {
+ return MultiValueResult.of(_segmentColumnReader.getDoubleMV(docId), null);
+ }
+
+ @Override
+ public String[] getStringMV(int docId) {
+ return _segmentColumnReader.getStringMV(docId);
+ }
+
+ @Override
+ public byte[][] getBytesMV(int docId) {
+ return _segmentColumnReader.getBytesMV(docId);
+ }
+
+ /**
+ * Validate that the document ID is within valid range.
+ *
+ * @param docId Document ID to validate
+ * @throws IndexOutOfBoundsException if docId is out of range
+ */
+ private void validateDocId(int docId) {
+ if (docId < 0 || docId >= _numDocs) {
+ throw new IndexOutOfBoundsException(
+ "docId " + docId + " is out of range. Valid range is 0 to " + (_numDocs - 1));
+ }
+ }
+
+ @Override
+ public void close()
+ throws IOException {
+ _segmentColumnReader.close();
+ }
+}
diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnarDataSource.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnarDataSource.java
new file mode 100644
index 0000000000..960dc7ebc4
--- /dev/null
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnarDataSource.java
@@ -0,0 +1,63 @@
+/**
+ * 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.io.IOException;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.spi.data.readers.ColumnReaderFactory;
+import org.apache.pinot.spi.data.readers.ColumnarDataSource;
+
+
+/**
+ * ColumnarDataSource implementation that wraps a Pinot segment.
+ * Provides columnar access to segment data via ColumnReaderFactory.
+ */
+public class PinotSegmentColumnarDataSource implements ColumnarDataSource {
+
+ private final IndexSegment _indexSegment;
+ private final int _totalDocs;
+ private final boolean _initializeDefaultValueReaders;
+
+ public PinotSegmentColumnarDataSource(IndexSegment indexSegment, boolean initializeDefaultValueReaders) {
+ _indexSegment = indexSegment;
+ _totalDocs = indexSegment.getSegmentMetadata().getTotalDocs();
+ _initializeDefaultValueReaders = initializeDefaultValueReaders;
+ }
+
+ @Override
+ public int getTotalDocs() {
+ return _totalDocs;
+ }
+
+ @Override
+ public ColumnReaderFactory createColumnReaderFactory() {
+ return new PinotSegmentColumnReaderFactory(_indexSegment, _initializeDefaultValueReaders);
+ }
+
+ @Override
+ public void close()
+ throws IOException {
+ // Segment lifecycle is managed externally, so no cleanup needed here
+ }
+
+ @Override
+ public String toString() {
+ return "PinotSegmentColumnarDataSource{segment=" + _indexSegment.getSegmentName() + "}";
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnReaderInterfaceTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnReaderInterfaceTest.java
new file mode 100644
index 0000000000..42f280e170
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnReaderInterfaceTest.java
@@ -0,0 +1,129 @@
+/**
+ * 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.creator.impl;
+
+import java.io.File;
+import java.util.Map;
+import java.util.Set;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReaderFactory;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * Tests for core ColumnReader interface and implementations.
+ *
+ * This test class validates:
+ *
+ * - PinotSegmentColumnReaderFactory functionality
+ * - DefaultValueColumnReader for new columns
+ * - ColumnReader interface contract
+ *
+ */
+public class ColumnReaderInterfaceTest extends ColumnarSegmentBuildingTestBase {
+
+ @Test
+ public void testColumnReaderFactory()
+ throws Exception {
+ // Create a segment to test column reader factory with
+ File segmentDir = createRowMajorSegment();
+ ImmutableSegment segment = ImmutableSegmentLoader.load(segmentDir, ReadMode.mmap);
+
+ try {
+ // Test PinotSegmentColumnReaderFactory
+ try (PinotSegmentColumnReaderFactory factory = new PinotSegmentColumnReaderFactory(segment)) {
+ factory.init(_originalSchema);
+
+ // Test that all expected columns are available
+ Set availableColumns = factory.getAvailableColumns();
+ Assert.assertTrue(availableColumns.contains(STRING_COL_1));
+ Assert.assertTrue(availableColumns.contains(INT_COL_1));
+ Assert.assertTrue(availableColumns.contains(MV_INT_COL));
+
+ // Test getting individual column readers
+ ColumnReader stringReader = factory.getColumnReader(STRING_COL_1);
+ Assert.assertEquals(stringReader.getColumnName(), STRING_COL_1);
+
+ ColumnReader mvReader = factory.getColumnReader(MV_INT_COL);
+ Assert.assertEquals(mvReader.getColumnName(), MV_INT_COL);
+
+ // Test reading values using iterator pattern
+ Assert.assertTrue(stringReader.hasNext());
+ Object firstStringValue = stringReader.next();
+ Assert.assertEquals(firstStringValue, "string1_0");
+
+ // Test that we can continue reading
+ Assert.assertTrue(stringReader.hasNext());
+ Object secondStringValue = stringReader.next();
+ Assert.assertEquals(secondStringValue, "string1_1");
+
+ // Reset and test again
+ stringReader.rewind();
+ Assert.assertTrue(stringReader.hasNext());
+ firstStringValue = stringReader.next();
+ Assert.assertEquals(firstStringValue, "string1_0");
+
+ // Test getting all column readers
+ Map allReaders = factory.getAllColumnReaders();
+ Assert.assertEquals(allReaders.size(), _originalSchema.getPhysicalColumnNames().size());
+ }
+ } finally {
+ segment.destroy();
+ }
+ }
+
+ @Test
+ public void testColumnReaderWithNewColumns()
+ throws Exception {
+ // Create a segment to test with
+ File segmentDir = createRowMajorSegment();
+ ImmutableSegment segment = ImmutableSegmentLoader.load(segmentDir, ReadMode.mmap);
+
+ try {
+ // Test getting readers for new columns (should return default value readers)
+ try (PinotSegmentColumnReaderFactory factory = new PinotSegmentColumnReaderFactory(segment)) {
+ factory.init(_extendedSchema);
+
+ // Test getting reader for new column
+ ColumnReader newStringReader = factory.getColumnReader(NEW_STRING_COL);
+ Assert.assertEquals(newStringReader.getColumnName(), NEW_STRING_COL);
+
+ // Verify it returns default values using iterator pattern
+ Assert.assertTrue(newStringReader.hasNext());
+ Object defaultValue = newStringReader.next();
+ Assert.assertEquals(defaultValue, _extendedSchema.getFieldSpecFor(NEW_STRING_COL).getDefaultNullValue());
+
+ // Test that all values are the same (default)
+ int valueCount = 1; // We already read one value
+ while (newStringReader.hasNext()) {
+ Object value = newStringReader.next();
+ Assert.assertEquals(value, defaultValue);
+ valueCount++;
+ }
+ Assert.assertEquals(valueCount, _testData.size());
+ }
+ } finally {
+ segment.destroy();
+ }
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarRowMajorEquivalenceTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarRowMajorEquivalenceTest.java
new file mode 100644
index 0000000000..ac2cac33f7
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarRowMajorEquivalenceTest.java
@@ -0,0 +1,85 @@
+/**
+ * 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.creator.impl;
+
+import java.io.File;
+import java.util.Set;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * Tests for columnar vs row-major equivalence.
+ *
+ * This test class validates that:
+ *
+ * - Columnar segment building produces identical results to row-major building
+ * - All data types are correctly handled in columnar mode
+ * - Segment statistics and metadata are preserved
+ *
+ */
+public class ColumnarRowMajorEquivalenceTest extends ColumnarSegmentBuildingTestBase {
+
+ @Test
+ public void testBasicColumnarBuilding()
+ throws Exception {
+ // First create a segment using traditional row-major approach
+ File rowMajorSegmentDir = createRowMajorSegment();
+
+ // Then create a segment using columnar approach from the row-major segment
+ File columnarSegmentDir = createColumnarSegment(rowMajorSegmentDir);
+
+ // Validate that both segments have identical data
+ validateSegmentsIdentical(rowMajorSegmentDir, columnarSegmentDir);
+ }
+
+ @Test
+ public void testAllDataTypes()
+ throws Exception {
+ // This test validates that all supported data types work correctly with columnar building
+ File rowMajorSegmentDir = createRowMajorSegment();
+ File columnarSegmentDir = createColumnarSegment(rowMajorSegmentDir);
+
+ // Validate that both segments have identical data for all data types
+ validateSegmentsIdentical(rowMajorSegmentDir, columnarSegmentDir);
+
+ // Additionally validate that all expected columns and data types are present
+ ImmutableSegment segment = ImmutableSegmentLoader.load(columnarSegmentDir, ReadMode.mmap);
+ try {
+ Set columnNames = segment.getPhysicalColumnNames();
+
+ // Validate all data types are present
+ Assert.assertTrue(columnNames.contains(STRING_COL_1), "STRING column missing");
+ Assert.assertTrue(columnNames.contains(INT_COL_1), "INT column missing");
+ Assert.assertTrue(columnNames.contains(LONG_COL), "LONG column missing");
+ Assert.assertTrue(columnNames.contains(FLOAT_COL), "FLOAT column missing");
+ Assert.assertTrue(columnNames.contains(DOUBLE_COL), "DOUBLE column missing");
+ Assert.assertTrue(columnNames.contains(BIG_DECIMAL_COL), "BIG_DECIMAL column missing");
+ Assert.assertTrue(columnNames.contains(BYTES_COL), "BYTES column missing");
+ Assert.assertTrue(columnNames.contains(MV_INT_COL), "Multi-value INT column missing");
+ Assert.assertTrue(columnNames.contains(MV_STRING_COL), "Multi-value STRING column missing");
+ Assert.assertTrue(columnNames.contains(TIME_COL), "TIME column missing");
+ } finally {
+ segment.destroy();
+ }
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSchemaEvolutionTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSchemaEvolutionTest.java
new file mode 100644
index 0000000000..ab533bb32c
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSchemaEvolutionTest.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.segment.creator.impl;
+
+import java.io.File;
+import org.testng.annotations.Test;
+
+
+/**
+ * Tests for schema evolution and new column handling.
+ *
+ * This test class validates:
+ *
+ * - Adding new columns with default values during columnar building
+ * - Schema evolution compatibility between row-major and columnar approaches
+ * - Default value handling for single-value and multi-value columns
+ *
+ */
+public class ColumnarSchemaEvolutionTest extends ColumnarSegmentBuildingTestBase {
+
+ @Test
+ public void testColumnarBuildingWithNewColumns()
+ throws Exception {
+ // Create original segment with original schema
+ File originalSegmentDir = createRowMajorSegment();
+
+ // Create columnar segment with extended schema (has additional columns)
+ File columnarSegmentDir = createColumnarSegmentWithNewColumns(originalSegmentDir);
+
+ // Create row-major segment with extended schema for comparison using the same original segment
+ File rowMajorExtendedSegmentDir = createRowMajorSegmentWithExtendedSchema(originalSegmentDir);
+
+ // Validate that both segments are identical (comprehensive comparison)
+ validateSegmentsIdentical(rowMajorExtendedSegmentDir, columnarSegmentDir);
+
+ // Additional validation that the new segment has the additional columns with default values
+ validateSegmentWithNewColumns(columnarSegmentDir);
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSegmentBuildingTestBase.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSegmentBuildingTestBase.java
new file mode 100644
index 0000000000..82dd918124
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSegmentBuildingTestBase.java
@@ -0,0 +1,578 @@
+/**
+ * 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.creator.impl;
+
+import com.google.common.collect.Lists;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReaderFactory;
+import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader;
+import org.apache.pinot.segment.spi.ColumnMetadata;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.RecordReader;
+import org.apache.pinot.spi.data.readers.RecordReaderConfig;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+
+
+/**
+ * Base class for columnar segment building tests containing common setup and utility methods.
+ *
+ * This base class provides:
+ *
+ * - Common test data and schema setup
+ * - Segment creation utility methods
+ * - Validation utility methods
+ * - Test data generation
+ *
+ */
+public abstract class ColumnarSegmentBuildingTestBase {
+ protected static final String TEMP_DIR = System.getProperty("java.io.tmpdir");
+ protected static final String TABLE_NAME = "testTable";
+ protected static final String SEGMENT_NAME = "testSegment";
+
+ // Test columns
+ protected static final String STRING_COL_1 = "stringCol1";
+ protected static final String STRING_COL_2 = "stringCol2";
+ protected static final String INT_COL_1 = "intCol1";
+ protected static final String INT_COL_2 = "intCol2";
+ protected static final String LONG_COL = "longCol";
+ protected static final String FLOAT_COL = "floatCol";
+ protected static final String DOUBLE_COL = "doubleCol";
+ protected static final String BIG_DECIMAL_COL = "bigDecimalCol";
+ protected static final String BYTES_COL = "bytesCol";
+ protected static final String TIME_COL = "timeCol";
+ protected static final String MV_INT_COL = "mvIntCol";
+ protected static final String MV_LONG_COL = "mvLongCol";
+ protected static final String MV_FLOAT_COL = "mvFloatCol";
+ protected static final String MV_DOUBLE_COL = "mvDoubleCol";
+ protected static final String MV_STRING_COL = "mvStringCol";
+ protected static final String MV_BYTES_COL = "mvBytesCol";
+
+ // New column for testing default value handling
+ protected static final String NEW_STRING_COL = "newStringCol";
+ protected static final String NEW_INT_COL = "newIntCol";
+ protected static final String NEW_MV_LONG_COL = "newMvLongCol";
+
+ protected File _tempDir;
+ protected Schema _originalSchema;
+ protected Schema _extendedSchema; // Schema with additional columns
+ protected TableConfig _tableConfig;
+ protected List _testData;
+
+ @BeforeClass
+ public void setUp()
+ throws IOException {
+ _tempDir = new File(TEMP_DIR, getClass().getSimpleName());
+ FileUtils.deleteQuietly(_tempDir);
+ _tempDir.mkdirs();
+
+ // Create original schema
+ _originalSchema = new Schema.SchemaBuilder()
+ .addSingleValueDimension(STRING_COL_1, FieldSpec.DataType.STRING)
+ .addSingleValueDimension(STRING_COL_2, FieldSpec.DataType.STRING)
+ .addSingleValueDimension(INT_COL_1, FieldSpec.DataType.INT)
+ .addSingleValueDimension(INT_COL_2, FieldSpec.DataType.INT)
+ .addSingleValueDimension(LONG_COL, FieldSpec.DataType.LONG)
+ .addSingleValueDimension(FLOAT_COL, FieldSpec.DataType.FLOAT)
+ .addSingleValueDimension(DOUBLE_COL, FieldSpec.DataType.DOUBLE)
+ .addSingleValueDimension(BIG_DECIMAL_COL, FieldSpec.DataType.BIG_DECIMAL)
+ .addSingleValueDimension(BYTES_COL, FieldSpec.DataType.BYTES)
+ .addMultiValueDimension(MV_INT_COL, FieldSpec.DataType.INT)
+ .addMultiValueDimension(MV_LONG_COL, FieldSpec.DataType.LONG)
+ .addMultiValueDimension(MV_FLOAT_COL, FieldSpec.DataType.FLOAT)
+ .addMultiValueDimension(MV_DOUBLE_COL, FieldSpec.DataType.DOUBLE)
+ .addMultiValueDimension(MV_STRING_COL, FieldSpec.DataType.STRING)
+ .addMultiValueDimension(MV_BYTES_COL, FieldSpec.DataType.BYTES)
+ .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .build();
+
+ // Create extended schema with additional columns
+ _extendedSchema = new Schema.SchemaBuilder()
+ .addSingleValueDimension(STRING_COL_1, FieldSpec.DataType.STRING)
+ .addSingleValueDimension(STRING_COL_2, FieldSpec.DataType.STRING)
+ .addSingleValueDimension(INT_COL_1, FieldSpec.DataType.INT)
+ .addSingleValueDimension(INT_COL_2, FieldSpec.DataType.INT)
+ .addSingleValueDimension(LONG_COL, FieldSpec.DataType.LONG)
+ .addSingleValueDimension(FLOAT_COL, FieldSpec.DataType.FLOAT)
+ .addSingleValueDimension(DOUBLE_COL, FieldSpec.DataType.DOUBLE)
+ .addSingleValueDimension(BIG_DECIMAL_COL, FieldSpec.DataType.BIG_DECIMAL)
+ .addSingleValueDimension(BYTES_COL, FieldSpec.DataType.BYTES)
+ .addMultiValueDimension(MV_INT_COL, FieldSpec.DataType.INT)
+ .addMultiValueDimension(MV_LONG_COL, FieldSpec.DataType.LONG)
+ .addMultiValueDimension(MV_FLOAT_COL, FieldSpec.DataType.FLOAT)
+ .addMultiValueDimension(MV_DOUBLE_COL, FieldSpec.DataType.DOUBLE)
+ .addMultiValueDimension(MV_STRING_COL, FieldSpec.DataType.STRING)
+ .addMultiValueDimension(MV_BYTES_COL, FieldSpec.DataType.BYTES)
+ .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .addSingleValueDimension(NEW_STRING_COL, FieldSpec.DataType.STRING)
+ .addSingleValueDimension(NEW_INT_COL, FieldSpec.DataType.INT)
+ .addMultiValueDimension(NEW_MV_LONG_COL, FieldSpec.DataType.LONG)
+ .build();
+
+ // Create table config
+ _tableConfig = new TableConfigBuilder(TableType.OFFLINE)
+ .setTableName(TABLE_NAME)
+ .setTimeColumnName(TIME_COL)
+ .setInvertedIndexColumns(Lists.newArrayList(STRING_COL_1, INT_COL_1))
+ .setSortedColumn(INT_COL_1)
+ .build();
+
+ // Generate test data
+ _testData = generateTestData(100);
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws IOException {
+ FileUtils.deleteQuietly(_tempDir);
+ }
+
+ protected File createRowMajorSegment()
+ throws Exception {
+ return createRowMajorSegmentWithConfig(_tableConfig, _originalSchema, "rowMajorSegment");
+ }
+
+ protected File createRowMajorSegmentWithConfig(TableConfig tableConfig, Schema schema, String dirName)
+ throws Exception {
+ File outputDir = new File(_tempDir, dirName);
+ FileUtils.deleteQuietly(outputDir);
+ outputDir.mkdirs();
+
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema);
+ config.setOutDir(outputDir.getAbsolutePath());
+ config.setSegmentName(SEGMENT_NAME + "_" + dirName);
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+ TestRecordReader recordReader = new TestRecordReader(_testData);
+ driver.init(config, recordReader);
+ driver.build();
+
+ return new File(outputDir, SEGMENT_NAME + "_" + dirName);
+ }
+
+ protected TableConfig createTableConfigWithNoDictionary() {
+ List allColumns = new ArrayList<>();
+ allColumns.add(STRING_COL_1);
+ allColumns.add(STRING_COL_2);
+ allColumns.add(INT_COL_1);
+ allColumns.add(INT_COL_2);
+ allColumns.add(LONG_COL);
+ allColumns.add(FLOAT_COL);
+ allColumns.add(DOUBLE_COL);
+ allColumns.add(BIG_DECIMAL_COL);
+ allColumns.add(BYTES_COL);
+ allColumns.add(MV_INT_COL);
+ allColumns.add(MV_LONG_COL);
+ allColumns.add(MV_FLOAT_COL);
+ allColumns.add(MV_DOUBLE_COL);
+ allColumns.add(MV_STRING_COL);
+ allColumns.add(MV_BYTES_COL);
+
+ return new TableConfigBuilder(TableType.OFFLINE)
+ .setTableName(TABLE_NAME)
+ .setTimeColumnName(TIME_COL)
+ .setNoDictionaryColumns(allColumns)
+ .build();
+ }
+
+ protected Schema createNullableSchema() {
+ return new Schema.SchemaBuilder()
+ .setEnableColumnBasedNullHandling(true)
+ .addDimensionField(STRING_COL_1, FieldSpec.DataType.STRING, spec -> spec.setNullable(true))
+ .addDimensionField(STRING_COL_2, FieldSpec.DataType.STRING, spec -> spec.setNullable(true))
+ .addDimensionField(INT_COL_1, FieldSpec.DataType.INT, spec -> spec.setNullable(true))
+ .addDimensionField(INT_COL_2, FieldSpec.DataType.INT, spec -> spec.setNullable(true))
+ .addDimensionField(LONG_COL, FieldSpec.DataType.LONG, spec -> spec.setNullable(true))
+ .addDimensionField(FLOAT_COL, FieldSpec.DataType.FLOAT, spec -> spec.setNullable(true))
+ .addDimensionField(DOUBLE_COL, FieldSpec.DataType.DOUBLE, spec -> spec.setNullable(true))
+ .addDimensionField(BIG_DECIMAL_COL, FieldSpec.DataType.BIG_DECIMAL, spec -> spec.setNullable(true))
+ .addDimensionField(BYTES_COL, FieldSpec.DataType.BYTES, spec -> spec.setNullable(true))
+ .addDimensionField(MV_INT_COL, FieldSpec.DataType.INT, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDimensionField(MV_LONG_COL, FieldSpec.DataType.LONG, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDimensionField(MV_FLOAT_COL, FieldSpec.DataType.FLOAT, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDimensionField(MV_DOUBLE_COL, FieldSpec.DataType.DOUBLE, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDimensionField(MV_STRING_COL, FieldSpec.DataType.STRING, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDimensionField(MV_BYTES_COL, FieldSpec.DataType.BYTES, spec -> {
+ spec.setSingleValueField(false);
+ spec.setNullable(true);
+ })
+ .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .build();
+ }
+
+ protected File createColumnarSegment(File sourceSegmentDir)
+ throws Exception {
+ File outputDir = new File(_tempDir, "columnarSegment");
+ FileUtils.deleteQuietly(outputDir);
+ outputDir.mkdirs();
+
+ // Load the source segment
+ ImmutableSegment sourceSegment = ImmutableSegmentLoader.load(sourceSegmentDir, ReadMode.mmap);
+
+ try {
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(_tableConfig, _originalSchema);
+ config.setOutDir(outputDir.getAbsolutePath());
+ config.setSegmentName(SEGMENT_NAME + "_columnar");
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+
+ // Use the new columnar building approach with config
+ try (PinotSegmentColumnReaderFactory factory = new PinotSegmentColumnReaderFactory(sourceSegment)) {
+ driver.init(config, factory);
+ driver.build();
+ }
+
+ return new File(outputDir, SEGMENT_NAME + "_columnar");
+ } finally {
+ sourceSegment.destroy();
+ }
+ }
+
+ protected File createColumnarSegmentWithNewColumns(File sourceSegmentDir)
+ throws Exception {
+ File outputDir = new File(_tempDir, "columnarSegmentWithNewColumns");
+ FileUtils.deleteQuietly(outputDir);
+ outputDir.mkdirs();
+
+ // Load the source segment
+ ImmutableSegment sourceSegment = ImmutableSegmentLoader.load(sourceSegmentDir, ReadMode.mmap);
+
+ try {
+ // Use extended schema with new columns
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(_tableConfig, _extendedSchema);
+ config.setOutDir(outputDir.getAbsolutePath());
+ config.setSegmentName(SEGMENT_NAME + "_withNewColumns");
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+
+ // Use columnar building with extended schema
+ try (PinotSegmentColumnReaderFactory factory = new PinotSegmentColumnReaderFactory(sourceSegment)) {
+ driver.init(config, factory);
+ driver.build();
+ }
+
+ return new File(outputDir, SEGMENT_NAME + "_withNewColumns");
+ } finally {
+ sourceSegment.destroy();
+ }
+ }
+
+ protected File createRowMajorSegmentWithExtendedSchema(File originalSegmentDir)
+ throws Exception {
+ File outputDir = new File(_tempDir, "rowMajorSegmentWithExtendedSchema");
+ FileUtils.deleteQuietly(outputDir);
+ outputDir.mkdirs();
+
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(_tableConfig, _extendedSchema);
+ config.setOutDir(outputDir.getAbsolutePath());
+ config.setSegmentName(SEGMENT_NAME + "_rowMajorExtended");
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+
+ // Use PinotSegmentRecordReader to read from the original segment, similar to RefreshSegmentTaskExecutor
+ try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) {
+ recordReader.init(originalSegmentDir, null, null);
+ driver.init(config, recordReader);
+ driver.build();
+ }
+
+ return new File(outputDir, SEGMENT_NAME + "_rowMajorExtended");
+ }
+
+ protected void validateSegmentsIdentical(File segment1Dir, File segment2Dir)
+ throws Exception {
+ ImmutableSegment segment1 = ImmutableSegmentLoader.load(segment1Dir, ReadMode.mmap);
+ ImmutableSegment segment2 = ImmutableSegmentLoader.load(segment2Dir, ReadMode.mmap);
+
+ try {
+ // Validate basic segment metadata
+ Assert.assertEquals(segment1.getSegmentMetadata().getTotalDocs(), segment2.getSegmentMetadata().getTotalDocs());
+ Assert.assertEquals(segment1.getSegmentMetadata().getAllColumns(), segment2.getSegmentMetadata().getAllColumns());
+
+ // Validate column metadata and statistics for each column
+ for (String columnName : segment1.getPhysicalColumnNames()) {
+ validateColumnMetadata(segment1, segment2, columnName);
+ validateColumnData(segment1, segment2, columnName);
+ }
+ } finally {
+ segment1.destroy();
+ segment2.destroy();
+ }
+ }
+
+ protected void validateColumnMetadata(ImmutableSegment segment1, ImmutableSegment segment2, String columnName) {
+ ColumnMetadata metadata1 = segment1.getSegmentMetadata().getColumnMetadataFor(columnName);
+ ColumnMetadata metadata2 = segment2.getSegmentMetadata().getColumnMetadataFor(columnName);
+
+ Assert.assertNotNull(metadata1, "Column metadata missing for " + columnName + " in segment1");
+ Assert.assertNotNull(metadata2, "Column metadata missing for " + columnName + " in segment2");
+
+ // Validate basic column properties
+ Assert.assertEquals(metadata1.getDataType(), metadata2.getDataType(),
+ "Data type mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.isSingleValue(), metadata2.isSingleValue(),
+ "Single value flag mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.getFieldType(), metadata2.getFieldType(),
+ "Field type mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.getTotalDocs(), metadata2.getTotalDocs(),
+ "Total docs mismatch for column " + columnName);
+
+ // Validate cardinality and dictionary properties
+ Assert.assertEquals(metadata1.getCardinality(), metadata2.getCardinality(),
+ "Cardinality mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.hasDictionary(), metadata2.hasDictionary(),
+ "Dictionary flag mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.getColumnMaxLength(), metadata2.getColumnMaxLength(),
+ "Column max length mismatch for column " + columnName);
+
+ // Validate sorting and indexing properties
+ Assert.assertEquals(metadata1.isSorted(), metadata2.isSorted(),
+ "Sorted flag mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.isAutoGenerated(), metadata2.isAutoGenerated(),
+ "Auto-generated flag mismatch for column " + columnName);
+
+ // Validate multi-value properties
+ if (!metadata1.isSingleValue()) {
+ Assert.assertEquals(metadata1.getMaxNumberOfMultiValues(), metadata2.getMaxNumberOfMultiValues(),
+ "Max number of multi-values mismatch for column " + columnName);
+ Assert.assertEquals(metadata1.getTotalNumberOfEntries(), metadata2.getTotalNumberOfEntries(),
+ "Total number of entries mismatch for column " + columnName);
+ }
+
+ // Validate min/max values if available
+ Comparable minValue1 = metadata1.getMinValue();
+ Comparable minValue2 = metadata2.getMinValue();
+ Comparable maxValue1 = metadata1.getMaxValue();
+ Comparable maxValue2 = metadata2.getMaxValue();
+
+ if (minValue1 != null && minValue2 != null) {
+ Assert.assertEquals(minValue1, minValue2,
+ "Min value mismatch for column " + columnName);
+ }
+ if (maxValue1 != null && maxValue2 != null) {
+ Assert.assertEquals(maxValue1, maxValue2,
+ "Max value mismatch for column " + columnName);
+ }
+
+ // Validate bits per element for dictionary encoded columns
+ if (metadata1.hasDictionary() && metadata2.hasDictionary()) {
+ Assert.assertEquals(metadata1.getBitsPerElement(), metadata2.getBitsPerElement(),
+ "Bits per element mismatch for column " + columnName);
+ }
+ }
+
+ protected void validateSegmentWithNewColumns(File segmentDir)
+ throws Exception {
+ ImmutableSegment segment = ImmutableSegmentLoader.load(segmentDir, ReadMode.mmap);
+
+ try {
+ // Validate that new columns exist
+ Assert.assertTrue(segment.getPhysicalColumnNames().contains(NEW_STRING_COL));
+ Assert.assertTrue(segment.getPhysicalColumnNames().contains(NEW_INT_COL));
+ Assert.assertTrue(segment.getPhysicalColumnNames().contains(NEW_MV_LONG_COL));
+
+ // Validate that new columns have default values
+ GenericRow row = new GenericRow();
+ for (int docId = 0; docId < segment.getSegmentMetadata().getTotalDocs(); docId++) {
+ segment.getRecord(docId, row);
+
+ // Check default values
+ Assert.assertEquals(row.getValue(NEW_STRING_COL),
+ _extendedSchema.getFieldSpecFor(NEW_STRING_COL).getDefaultNullValue());
+ Assert.assertEquals(row.getValue(NEW_INT_COL),
+ _extendedSchema.getFieldSpecFor(NEW_INT_COL).getDefaultNullValue());
+
+ // For multi-value columns, the default value is wrapped in an array
+ Object mvLongValue = row.getValue(NEW_MV_LONG_COL);
+ Assert.assertTrue(mvLongValue instanceof Object[], "Multi-value column should return an array");
+ Object[] mvLongArray = (Object[]) mvLongValue;
+ Assert.assertEquals(mvLongArray.length, 1, "Default multi-value array should have one element");
+ Assert.assertEquals(mvLongArray[0],
+ _extendedSchema.getFieldSpecFor(NEW_MV_LONG_COL).getDefaultNullValue());
+ }
+ } finally {
+ segment.destroy();
+ }
+ }
+
+ protected void validateColumnData(ImmutableSegment segment1, ImmutableSegment segment2, String columnName) {
+ int numDocs = segment1.getSegmentMetadata().getTotalDocs();
+
+ GenericRow row1 = new GenericRow();
+ GenericRow row2 = new GenericRow();
+
+ for (int docId = 0; docId < numDocs; docId++) {
+ segment1.getRecord(docId, row1);
+ segment2.getRecord(docId, row2);
+
+ Object value1 = row1.getValue(columnName);
+ Object value2 = row2.getValue(columnName);
+
+ Assert.assertEquals(value1, value2,
+ String.format("Column %s differs at docId %d: %s vs %s", columnName, docId, value1, value2));
+ }
+ }
+
+ protected List generateTestData(int numRows) {
+ List data = new ArrayList<>();
+
+ // Use a fixed seed for reproducible test results
+ Random random = new Random(42);
+
+ // Null probability - approximately 10% of values will be null for nullable columns
+ double nullProbability = 0.1;
+
+ for (int i = 0; i < numRows; i++) {
+ GenericRow row = new GenericRow();
+
+ // STRING columns - can be null
+ row.putValue(STRING_COL_1, random.nextDouble() < nullProbability ? null : "string1_" + i);
+ row.putValue(STRING_COL_2, random.nextDouble() < nullProbability ? null : "string2_" + (i % 10));
+
+ // INT columns - can be null
+ row.putValue(INT_COL_1, random.nextDouble() < nullProbability ? null : i);
+ row.putValue(INT_COL_2, random.nextDouble() < nullProbability ? null : i * 2);
+
+ // LONG column - can be null
+ row.putValue(LONG_COL, random.nextDouble() < nullProbability ? null : (long) i * 3);
+
+ // FLOAT column - can be null
+ row.putValue(FLOAT_COL, random.nextDouble() < nullProbability ? null : (float) i * 1.5);
+
+ // DOUBLE column - can be null
+ row.putValue(DOUBLE_COL, random.nextDouble() < nullProbability ? null : (double) i * 2.5);
+
+ // BIG_DECIMAL column - can be null
+ row.putValue(BIG_DECIMAL_COL, random.nextDouble() < nullProbability ? null : new BigDecimal(i + ".123"));
+
+ // BYTES column - can be null
+ row.putValue(BYTES_COL, random.nextDouble() < nullProbability ? null : ("bytes_" + i).getBytes());
+
+ // TIME column - typically not null for time columns, but we can test it
+ // Use a deterministic timestamp based on the row index for consistent results
+ long baseTimestamp = 1700000000000L; // Fixed base timestamp
+ row.putValue(TIME_COL, random.nextDouble() < nullProbability ? null : baseTimestamp + i);
+
+ // Multi-value columns - can be null (entire array is null, not individual elements)
+ row.putValue(MV_INT_COL, random.nextDouble() < nullProbability ? null : new Object[]{i, i + 1, i + 2});
+ row.putValue(MV_LONG_COL, random.nextDouble() < nullProbability ? null
+ : new Object[]{(long) i * 10, (long) (i + 1) * 10, (long) (i + 2) * 10});
+ row.putValue(MV_FLOAT_COL, random.nextDouble() < nullProbability ? null
+ : new Object[]{(float) i * 1.1f, (float) (i + 1) * 1.1f, (float) (i + 2) * 1.1f});
+ row.putValue(MV_DOUBLE_COL, random.nextDouble() < nullProbability ? null
+ : new Object[]{(double) i * 2.2, (double) (i + 1) * 2.2, (double) (i + 2) * 2.2});
+ row.putValue(MV_STRING_COL, random.nextDouble() < nullProbability ? null
+ : new Object[]{"mv1_" + i, "mv2_" + i, "mv3_" + (i % 3)});
+ row.putValue(MV_BYTES_COL, random.nextDouble() < nullProbability ? null
+ : new Object[]{("mvbytes1_" + i).getBytes(), ("mvbytes2_" + i).getBytes()});
+
+ data.add(row);
+ }
+
+ return data;
+ }
+
+ /**
+ * Simple test record reader for the test data.
+ */
+ protected static class TestRecordReader implements RecordReader {
+ private final List _data;
+ private int _currentIndex = 0;
+
+ public TestRecordReader(List data) {
+ _data = data;
+ }
+
+ @Override
+ public void init(File dataFile, @Nullable Set fieldsToRead,
+ @Nullable RecordReaderConfig recordReaderConfig) {
+ _currentIndex = 0;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return _currentIndex < _data.size();
+ }
+
+ @Override
+ public GenericRow next() {
+ return next(new GenericRow());
+ }
+
+ @Override
+ public GenericRow next(GenericRow reuse) {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more records");
+ }
+
+ GenericRow sourceRow = _data.get(_currentIndex++);
+ reuse.clear();
+
+ for (Map.Entry entry : sourceRow.getFieldToValueMap().entrySet()) {
+ reuse.putValue(entry.getKey(), entry.getValue());
+ }
+
+ return reuse;
+ }
+
+ @Override
+ public void rewind() {
+ _currentIndex = 0;
+ }
+
+ @Override
+ public void close() {
+ // No-op
+ }
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollectorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollectorTest.java
new file mode 100644
index 0000000000..3be1c48aaf
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollectorTest.java
@@ -0,0 +1,712 @@
+/**
+ * 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.creator.impl.stats;
+
+import java.math.BigDecimal;
+import java.util.Collections;
+import java.util.Set;
+import org.apache.pinot.segment.spi.creator.StatsCollectorConfig;
+import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Comprehensive tests for AbstractColumnStatisticsCollector implementations.
+ * Tests all data types with both primitive and Object collect methods,
+ * single and multi-value scenarios, edge cases, and partition functionality.
+ */
+public class AbstractColumnStatisticsCollectorTest {
+
+ private static final String COLUMN_NAME = "testColumn";
+ private static final String TABLE_NAME = "testTable";
+ private static final int NUM_PARTITIONS = 4;
+
+ // Test helper methods
+
+ /**
+ * Creates a StatsCollectorConfig with optional partitioning.
+ */
+ private StatsCollectorConfig createStatsCollectorConfig(FieldSpec.DataType dataType, boolean isSingleValue,
+ boolean withPartitioning) {
+ TableConfigBuilder builder = new TableConfigBuilder(org.apache.pinot.spi.config.table.TableType.OFFLINE)
+ .setTableName(TABLE_NAME);
+
+ if (withPartitioning) {
+ builder.setSegmentPartitionConfig(new SegmentPartitionConfig(
+ Collections.singletonMap(COLUMN_NAME, new ColumnPartitionConfig("Murmur", NUM_PARTITIONS))));
+ }
+
+ TableConfig tableConfig = builder.build();
+ Schema schema = new Schema();
+ schema.addField(new DimensionFieldSpec(COLUMN_NAME, dataType, isSingleValue));
+
+ return new StatsCollectorConfig(tableConfig, schema,
+ withPartitioning ? tableConfig.getIndexingConfig().getSegmentPartitionConfig() : null);
+ }
+
+ /**
+ * Factory method to create the appropriate collector for a given data type.
+ */
+ private AbstractColumnStatisticsCollector createCollector(FieldSpec.DataType dataType, boolean isSingleValue,
+ boolean withPartitioning) {
+ StatsCollectorConfig config = createStatsCollectorConfig(dataType, isSingleValue, withPartitioning);
+ switch (dataType) {
+ case INT:
+ return new IntColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ case LONG:
+ return new LongColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ case FLOAT:
+ return new FloatColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ case DOUBLE:
+ return new DoubleColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ case STRING:
+ return new StringColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ case BYTES:
+ return new BytesColumnPredIndexStatsCollector(COLUMN_NAME, config);
+ case BIG_DECIMAL:
+ return new BigDecimalColumnPreIndexStatsCollector(COLUMN_NAME, config);
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ /**
+ * Helper method to assert basic statistics after collection.
+ */
+ private void assertBasicStats(AbstractColumnStatisticsCollector collector, int expectedCardinality,
+ int expectedTotalEntries, Object expectedMin, Object expectedMax, boolean expectedSorted) {
+ assertEquals(collector.getCardinality(), expectedCardinality);
+ assertEquals(collector.getTotalNumberOfEntries(), expectedTotalEntries);
+ assertEquals(collector.isSorted(), expectedSorted);
+
+ if (expectedMin != null && expectedMax != null) {
+ assertEquals(collector.getMinValue(), expectedMin);
+ assertEquals(collector.getMaxValue(), expectedMax);
+ }
+ }
+
+ // DataProviders
+
+ @DataProvider(name = "allDataTypes")
+ public Object[][] allDataTypes() {
+ return new Object[][] {
+ {FieldSpec.DataType.INT},
+ {FieldSpec.DataType.LONG},
+ {FieldSpec.DataType.FLOAT},
+ {FieldSpec.DataType.DOUBLE},
+ {FieldSpec.DataType.STRING},
+ {FieldSpec.DataType.BYTES},
+ {FieldSpec.DataType.BIG_DECIMAL}
+ };
+ }
+
+ @DataProvider(name = "primitiveDataTypes")
+ public Object[][] primitiveDataTypes() {
+ return new Object[][] {
+ {FieldSpec.DataType.INT},
+ {FieldSpec.DataType.LONG},
+ {FieldSpec.DataType.FLOAT},
+ {FieldSpec.DataType.DOUBLE}
+ };
+ }
+
+ @DataProvider(name = "numericDataTypes")
+ public Object[][] numericDataTypes() {
+ return new Object[][] {
+ {FieldSpec.DataType.INT},
+ {FieldSpec.DataType.LONG},
+ {FieldSpec.DataType.FLOAT},
+ {FieldSpec.DataType.DOUBLE},
+ {FieldSpec.DataType.BIG_DECIMAL}
+ };
+ }
+
+ @DataProvider(name = "multiValueSupportedTypes")
+ public Object[][] multiValueSupportedTypes() {
+ return new Object[][] {
+ {FieldSpec.DataType.INT},
+ {FieldSpec.DataType.LONG},
+ {FieldSpec.DataType.FLOAT},
+ {FieldSpec.DataType.DOUBLE},
+ {FieldSpec.DataType.STRING},
+ {FieldSpec.DataType.BYTES}
+ };
+ }
+
+ @DataProvider(name = "lengthTrackingTypes")
+ public Object[][] lengthTrackingTypes() {
+ return new Object[][] {
+ {FieldSpec.DataType.STRING},
+ {FieldSpec.DataType.BYTES},
+ {FieldSpec.DataType.BIG_DECIMAL}
+ };
+ }
+
+ // Test 1: Single-Value Collection using Primitive Methods
+
+ @Test(dataProvider = "primitiveDataTypes")
+ public void testSingleValuePrimitiveCollectionSorted(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect sorted values
+ switch (dataType) {
+ case INT:
+ collector.collect(1);
+ collector.collect(5);
+ collector.collect(5); // duplicate
+ collector.collect(10);
+ break;
+ case LONG:
+ collector.collect(1L);
+ collector.collect(5L);
+ collector.collect(5L); // duplicate
+ collector.collect(10L);
+ break;
+ case FLOAT:
+ collector.collect(1.5f);
+ collector.collect(5.5f);
+ collector.collect(5.5f); // duplicate
+ collector.collect(10.5f);
+ break;
+ case DOUBLE:
+ collector.collect(1.5);
+ collector.collect(5.5);
+ collector.collect(5.5); // duplicate
+ collector.collect(10.5);
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+
+ collector.seal();
+
+ // Verify stats
+ assertBasicStats(collector, 3, 4, getMinValue(dataType), getMaxValue(dataType), true);
+
+ // Verify unique values set
+ Object uniqueValuesSet = collector.getUniqueValuesSet();
+ assertNotNull(uniqueValuesSet);
+ assertEquals(getArrayLength(uniqueValuesSet), 3);
+ }
+
+ @Test(dataProvider = "primitiveDataTypes")
+ public void testSingleValuePrimitiveCollectionUnsorted(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect unsorted values
+ switch (dataType) {
+ case INT:
+ collector.collect(10);
+ collector.collect(1);
+ collector.collect(5);
+ collector.collect(5); // duplicate
+ break;
+ case LONG:
+ collector.collect(10L);
+ collector.collect(1L);
+ collector.collect(5L);
+ collector.collect(5L); // duplicate
+ break;
+ case FLOAT:
+ collector.collect(10.5f);
+ collector.collect(1.5f);
+ collector.collect(5.5f);
+ collector.collect(5.5f); // duplicate
+ break;
+ case DOUBLE:
+ collector.collect(10.5);
+ collector.collect(1.5);
+ collector.collect(5.5);
+ collector.collect(5.5); // duplicate
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+
+ collector.seal();
+
+ // Verify stats - should not be sorted
+ assertBasicStats(collector, 3, 4, getMinValue(dataType), getMaxValue(dataType), false);
+ }
+
+ // Test 2: Single-Value Collection using collect(Object)
+
+ @Test(dataProvider = "allDataTypes")
+ public void testSingleValueObjectCollection(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect values using Object method
+ Object[] testData = getTestDataForType(dataType, true);
+ for (Object value : testData) {
+ collector.collect(value);
+ }
+
+ collector.seal();
+
+ // Verify stats
+ int expectedCardinality = 3; // test data has 1 duplicate
+ int expectedTotalEntries = 4;
+ assertBasicStats(collector, expectedCardinality, expectedTotalEntries,
+ getExpectedMinForType(dataType), getExpectedMaxForType(dataType), true);
+ }
+
+ // Test 3: Multi-Value Collection using Primitive Arrays
+
+ @Test(dataProvider = "multiValueSupportedTypes")
+ public void testMultiValuePrimitiveArrayCollection(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, false, false);
+
+ // Collect multi-value entries
+ switch (dataType) {
+ case INT:
+ collector.collect(new int[]{1, 2});
+ collector.collect(new int[]{3, 4, 5});
+ break;
+ case LONG:
+ collector.collect(new long[]{1L, 2L});
+ collector.collect(new long[]{3L, 4L, 5L});
+ break;
+ case FLOAT:
+ collector.collect(new float[]{1.5f, 2.5f});
+ collector.collect(new float[]{3.5f, 4.5f, 5.5f});
+ break;
+ case DOUBLE:
+ collector.collect(new double[]{1.5, 2.5});
+ collector.collect(new double[]{3.5, 4.5, 5.5});
+ break;
+ case STRING:
+ collector.collect(new String[]{"a", "b"});
+ collector.collect(new String[]{"c", "d", "e"});
+ break;
+ case BYTES:
+ collector.collect(new byte[][]{{1}, {2}});
+ collector.collect(new byte[][]{{3}, {4}, {5}});
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+
+ collector.seal();
+
+ // Verify stats
+ assertEquals(collector.getCardinality(), 5);
+ assertEquals(collector.getTotalNumberOfEntries(), 5);
+ assertEquals(collector.getMaxNumberOfMultiValues(), 3);
+ assertFalse(collector.isSorted()); // multi-value is never sorted
+ }
+
+ @Test(dataProvider = "multiValueSupportedTypes")
+ public void testMultiValueObjectArrayCollection(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, false, false);
+
+ // Collect multi-value entries using Object[]
+ Object[][] testData = getMultiValueTestData(dataType);
+ for (Object[] values : testData) {
+ collector.collect((Object) values);
+ }
+
+ collector.seal();
+
+ // Verify stats
+ assertEquals(collector.getCardinality(), 5);
+ assertEquals(collector.getTotalNumberOfEntries(), 5);
+ assertEquals(collector.getMaxNumberOfMultiValues(), 3);
+ assertFalse(collector.isSorted());
+ }
+
+ // Test 4: Edge Cases
+
+ @Test(dataProvider = "allDataTypes")
+ public void testEmptyCollection(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Seal without collecting any data
+ collector.seal();
+
+ // Verify stats
+ assertEquals(collector.getCardinality(), 0);
+ assertEquals(collector.getTotalNumberOfEntries(), 0);
+ assertTrue(collector.isSorted()); // Empty is considered sorted
+ }
+
+ @Test(dataProvider = "allDataTypes")
+ public void testSingleValue(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect single value
+ Object singleValue = getSingleValueForType(dataType);
+ collector.collect(singleValue);
+ collector.seal();
+
+ // Verify stats - for BYTES, need to use ByteArray for comparison
+ Object expectedValue = dataType == FieldSpec.DataType.BYTES
+ ? new ByteArray((byte[]) singleValue)
+ : singleValue;
+ assertBasicStats(collector, 1, 1, expectedValue, expectedValue, true);
+ }
+
+ @Test(dataProvider = "allDataTypes")
+ public void testAllDuplicates(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect same value multiple times
+ Object value = getSingleValueForType(dataType);
+ for (int i = 0; i < 10; i++) {
+ collector.collect(value);
+ }
+ collector.seal();
+
+ // Verify stats - for BYTES, need to use ByteArray for comparison
+ Object expectedValue = dataType == FieldSpec.DataType.BYTES
+ ? new ByteArray((byte[]) value)
+ : value;
+ assertBasicStats(collector, 1, 10, expectedValue, expectedValue, true);
+ }
+
+ @Test(dataProvider = "allDataTypes", expectedExceptions = IllegalStateException.class)
+ public void testUnsealedAccessThrowsException(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Try to access stats before sealing - should throw
+ collector.getMinValue();
+ }
+
+ // Test 5: Partition Function Tests
+
+ @Test(dataProvider = "allDataTypes")
+ public void testPartitionFunctionEnabled(FieldSpec.DataType dataType) {
+ if (dataType == FieldSpec.DataType.BIG_DECIMAL) {
+ // Skip BigDecimal as partition testing is less relevant
+ return;
+ }
+
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, true);
+
+ // Collect values
+ Object[] testData = getTestDataForType(dataType, true);
+ for (Object value : testData) {
+ collector.collect(value);
+ }
+ collector.seal();
+
+ // Verify partition function is present
+ assertNotNull(collector.getPartitionFunction());
+ assertEquals(collector.getNumPartitions(), NUM_PARTITIONS);
+
+ // Verify partitions are populated
+ Set partitions = collector.getPartitions();
+ assertNotNull(partitions);
+ assertTrue(partitions.size() > 0);
+ assertTrue(partitions.size() <= NUM_PARTITIONS);
+
+ // All partition values should be within range
+ for (Integer partition : partitions) {
+ assertTrue(partition >= 0 && partition < NUM_PARTITIONS);
+ }
+ }
+
+ @Test(dataProvider = "allDataTypes")
+ public void testPartitionFunctionDisabled(FieldSpec.DataType dataType) {
+ AbstractColumnStatisticsCollector collector = createCollector(dataType, true, false);
+
+ // Collect values
+ Object[] testData = getTestDataForType(dataType, true);
+ for (Object value : testData) {
+ collector.collect(value);
+ }
+ collector.seal();
+
+ // Verify partition function is not present
+ assertNull(collector.getPartitionFunction());
+ assertNull(collector.getPartitions());
+ assertEquals(collector.getNumPartitions(), -1);
+ }
+
+ // Test 6: Type-Specific Tests
+
+ @Test
+ public void testStringLengthTracking() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.STRING, true, false);
+
+ collector.collect("a"); // length 1
+ collector.collect("xyz"); // length 3
+ collector.collect("hello"); // length 5
+
+ collector.seal();
+
+ assertEquals(collector.getLengthOfLargestElement(), 5);
+ assertEquals(collector.getMaxRowLengthInBytes(), 5);
+ }
+
+ @Test
+ public void testStringMultiValueLengthTracking() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.STRING, false, false);
+
+ collector.collect(new String[]{"a", "xy"}); // row length: 1 + 2 = 3
+ collector.collect(new String[]{"hello", "world"}); // row length: 5 + 5 = 10
+
+ collector.seal();
+
+ assertEquals(collector.getLengthOfLargestElement(), 5);
+ assertEquals(collector.getMaxRowLengthInBytes(), 10);
+ }
+
+ @Test
+ public void testBytesLengthTracking() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.BYTES, true, false);
+
+ collector.collect(new byte[]{1}); // length 1
+ collector.collect(new byte[]{1, 2, 3}); // length 3
+ collector.collect(new byte[]{1, 2, 3, 4, 5}); // length 5
+
+ collector.seal();
+
+ assertEquals(collector.getLengthOfShortestElement(), 1);
+ assertEquals(collector.getLengthOfLargestElement(), 5);
+ assertEquals(collector.getMaxRowLengthInBytes(), 5);
+ }
+
+ @Test
+ public void testBytesMultiValueLengthTracking() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.BYTES, false, false);
+
+ collector.collect(new byte[][]{{1}, {1, 2}}); // row length: 1 + 2 = 3
+ collector.collect(new byte[][]{{1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}); // row length: 5 + 5 = 10
+
+ collector.seal();
+
+ assertEquals(collector.getLengthOfShortestElement(), 1);
+ assertEquals(collector.getLengthOfLargestElement(), 5);
+ assertEquals(collector.getMaxRowLengthInBytes(), 10);
+ }
+
+ @Test
+ public void testBigDecimalLengthTracking() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.BIG_DECIMAL, true, false);
+
+ collector.collect(new BigDecimal("1.0"));
+ collector.collect(new BigDecimal("123456.789"));
+ collector.collect(new BigDecimal("99.99"));
+
+ collector.seal();
+
+ // BigDecimal tracks length
+ assertTrue(collector.getLengthOfShortestElement() > 0);
+ assertTrue(collector.getLengthOfLargestElement() > 0);
+ assertTrue(collector.getLengthOfLargestElement() >= collector.getLengthOfShortestElement());
+ }
+
+ @Test(expectedExceptions = UnsupportedOperationException.class)
+ public void testBigDecimalMultiValueNotSupported() {
+ AbstractColumnStatisticsCollector collector = createCollector(FieldSpec.DataType.BIG_DECIMAL, true, false);
+
+ // BigDecimal does not support multi-value
+ collector.collect((Object) new Object[]{new BigDecimal("1.0"), new BigDecimal("2.0")});
+ }
+
+ // Helper methods for test data generation
+
+ private Object[] getTestDataForType(FieldSpec.DataType dataType, boolean sorted) {
+ switch (dataType) {
+ case INT:
+ return sorted ? new Object[]{1, 5, 5, 10} : new Object[]{10, 1, 5, 5};
+ case LONG:
+ return sorted ? new Object[]{1L, 5L, 5L, 10L} : new Object[]{10L, 1L, 5L, 5L};
+ case FLOAT:
+ return sorted ? new Object[]{1.5f, 5.5f, 5.5f, 10.5f} : new Object[]{10.5f, 1.5f, 5.5f, 5.5f};
+ case DOUBLE:
+ return sorted ? new Object[]{1.5, 5.5, 5.5, 10.5} : new Object[]{10.5, 1.5, 5.5, 5.5};
+ case STRING:
+ return sorted ? new Object[]{"a", "b", "b", "c"} : new Object[]{"c", "a", "b", "b"};
+ case BYTES:
+ return sorted ? new Object[]{
+ new byte[]{1}, new byte[]{2}, new byte[]{2}, new byte[]{3}
+ } : new Object[]{
+ new byte[]{3}, new byte[]{1}, new byte[]{2}, new byte[]{2}
+ };
+ case BIG_DECIMAL:
+ return sorted ? new Object[]{
+ new BigDecimal("1.0"), new BigDecimal("5.0"), new BigDecimal("5.0"), new BigDecimal("10.0")
+ } : new Object[]{
+ new BigDecimal("10.0"), new BigDecimal("1.0"), new BigDecimal("5.0"), new BigDecimal("5.0")
+ };
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object[][] getMultiValueTestData(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return new Object[][]{
+ {1, 2},
+ {3, 4, 5}
+ };
+ case LONG:
+ return new Object[][]{
+ {1L, 2L},
+ {3L, 4L, 5L}
+ };
+ case FLOAT:
+ return new Object[][]{
+ {1.5f, 2.5f},
+ {3.5f, 4.5f, 5.5f}
+ };
+ case DOUBLE:
+ return new Object[][]{
+ {1.5, 2.5},
+ {3.5, 4.5, 5.5}
+ };
+ case STRING:
+ return new Object[][]{
+ {"a", "b"},
+ {"c", "d", "e"}
+ };
+ case BYTES:
+ return new Object[][]{
+ {new byte[]{1}, new byte[]{2}},
+ {new byte[]{3}, new byte[]{4}, new byte[]{5}}
+ };
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object getSingleValueForType(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return 42;
+ case LONG:
+ return 42L;
+ case FLOAT:
+ return 42.5f;
+ case DOUBLE:
+ return 42.5;
+ case STRING:
+ return "test";
+ case BYTES:
+ return new byte[]{1, 2, 3};
+ case BIG_DECIMAL:
+ return new BigDecimal("42.5");
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object getMinValue(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return 1;
+ case LONG:
+ return 1L;
+ case FLOAT:
+ return 1.5f;
+ case DOUBLE:
+ return 1.5;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object getMaxValue(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return 10;
+ case LONG:
+ return 10L;
+ case FLOAT:
+ return 10.5f;
+ case DOUBLE:
+ return 10.5;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object getExpectedMinForType(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return 1;
+ case LONG:
+ return 1L;
+ case FLOAT:
+ return 1.5f;
+ case DOUBLE:
+ return 1.5;
+ case STRING:
+ return "a";
+ case BYTES:
+ return new ByteArray(new byte[]{1});
+ case BIG_DECIMAL:
+ return new BigDecimal("1.0");
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private Object getExpectedMaxForType(FieldSpec.DataType dataType) {
+ switch (dataType) {
+ case INT:
+ return 10;
+ case LONG:
+ return 10L;
+ case FLOAT:
+ return 10.5f;
+ case DOUBLE:
+ return 10.5;
+ case STRING:
+ return "c";
+ case BYTES:
+ return new ByteArray(new byte[]{3});
+ case BIG_DECIMAL:
+ return new BigDecimal("10.0");
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + dataType);
+ }
+ }
+
+ private int getArrayLength(Object array) {
+ if (array instanceof int[]) {
+ return ((int[]) array).length;
+ } else if (array instanceof long[]) {
+ return ((long[]) array).length;
+ } else if (array instanceof float[]) {
+ return ((float[]) array).length;
+ } else if (array instanceof double[]) {
+ return ((double[]) array).length;
+ } else if (array instanceof Object[]) {
+ return ((Object[]) array).length;
+ }
+ throw new IllegalArgumentException("Unsupported array type: " + array.getClass());
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReaderTest.java
new file mode 100644
index 0000000000..abe05a48af
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReaderTest.java
@@ -0,0 +1,619 @@
+/**
+ * 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.io.IOException;
+import java.util.Arrays;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.MetricFieldSpec;
+import org.apache.pinot.spi.data.readers.MultiValueResult;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * Comprehensive tests for DefaultValueColumnReader.
+ *
+ * This test validates:
+ *
+ * - Single-value fields for all data types (INT, LONG, FLOAT, DOUBLE, STRING, BYTES)
+ * - Multi-value fields for all data types
+ * - Sequential access methods (next*, hasNext, skipNext, isNextNull)
+ * - Random access methods (get*, isNull)
+ * - Type indicator methods (isInt, isLong, etc.)
+ * - Rewind functionality
+ * - Boundary conditions and exception handling
+ * - Array reuse optimization for multi-value fields
+ *
+ */
+public class DefaultValueColumnReaderTest {
+
+ private static final int NUM_DOCS = 100;
+ private static final String COLUMN_NAME = "testColumn";
+
+ // ========== Single-Value Field Tests ==========
+
+ @Test
+ public void testSingleValueIntColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test metadata
+ Assert.assertEquals(reader.getColumnName(), COLUMN_NAME);
+ Assert.assertEquals(reader.getTotalDocs(), NUM_DOCS);
+ Assert.assertTrue(reader.isInt());
+ Assert.assertFalse(reader.isLong());
+ Assert.assertFalse(reader.isString());
+
+ // Test sequential access with nextInt
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertTrue(reader.hasNext());
+ Assert.assertFalse(reader.isNextNull());
+ Assert.assertEquals(reader.nextInt(), expectedValue);
+ }
+ Assert.assertFalse(reader.hasNext());
+
+ // Test rewind and next()
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.next(), expectedValue);
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertFalse(reader.isNull(i));
+ Assert.assertEquals(reader.getInt(i), expectedValue);
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleValueLongColumn() throws IOException {
+ FieldSpec fieldSpec = new MetricFieldSpec(COLUMN_NAME, FieldSpec.DataType.LONG);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertFalse(reader.isInt());
+ Assert.assertTrue(reader.isLong());
+ Assert.assertFalse(reader.isString());
+
+ // Test sequential access with nextLong
+ long expectedValue = ((Number) fieldSpec.getDefaultNullValue()).longValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertTrue(reader.hasNext());
+ Assert.assertEquals(reader.nextLong(), expectedValue);
+ }
+
+ // Test random access
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.getLong(i), expectedValue);
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleValueFloatColumn() throws IOException {
+ FieldSpec fieldSpec = new MetricFieldSpec(COLUMN_NAME, FieldSpec.DataType.FLOAT);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertTrue(reader.isFloat());
+ Assert.assertFalse(reader.isDouble());
+
+ // Test sequential access with nextFloat
+ float expectedValue = ((Number) fieldSpec.getDefaultNullValue()).floatValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.nextFloat(), expectedValue);
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.getFloat(i), expectedValue);
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleValueDoubleColumn() throws IOException {
+ FieldSpec fieldSpec = new MetricFieldSpec(COLUMN_NAME, FieldSpec.DataType.DOUBLE);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertFalse(reader.isFloat());
+ Assert.assertTrue(reader.isDouble());
+
+ // Test sequential access with nextDouble
+ double expectedValue = ((Number) fieldSpec.getDefaultNullValue()).doubleValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.nextDouble(), expectedValue);
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.getDouble(i), expectedValue);
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleValueStringColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.STRING, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertTrue(reader.isString());
+ Assert.assertFalse(reader.isBytes());
+
+ // Test sequential access with nextString
+ String expectedValue = (String) fieldSpec.getDefaultNullValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.nextString(), expectedValue);
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertEquals(reader.getString(i), expectedValue);
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleValueBytesColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.BYTES, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertFalse(reader.isString());
+ Assert.assertTrue(reader.isBytes());
+
+ // Test sequential access with nextBytes
+ byte[] expectedValue = (byte[]) fieldSpec.getDefaultNullValue();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertTrue(Arrays.equals(reader.nextBytes(), expectedValue));
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertTrue(Arrays.equals(reader.getBytes(i), expectedValue));
+ }
+
+ reader.close();
+ }
+
+ // ========== Multi-Value Field Tests ==========
+
+ @Test
+ public void testMultiValueIntColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertTrue(reader.isInt());
+
+ // Test sequential access with nextIntMV
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+ int[] expectedArray = new int[]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertTrue(reader.hasNext());
+ MultiValueResult mvResult = reader.nextIntMV();
+ Assert.assertFalse(mvResult.hasNulls());
+ Assert.assertTrue(Arrays.equals(mvResult.getValues(), expectedArray));
+ }
+
+ // Test random access
+ reader.rewind();
+ for (int i = 0; i < NUM_DOCS; i++) {
+ MultiValueResult mvResult = reader.getIntMV(i);
+ Assert.assertFalse(mvResult.hasNulls());
+ Assert.assertTrue(Arrays.equals(mvResult.getValues(), expectedArray));
+ }
+
+ // Test that the same array instance is returned (optimization)
+ reader.rewind();
+ int[] firstCall = reader.getIntMV(0).getValues();
+ int[] secondCall = reader.getIntMV(1).getValues();
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ @Test
+ public void testMultiValueLongColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.LONG, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test type indicators
+ Assert.assertTrue(reader.isLong());
+
+ // Test sequential access with nextLongMV
+ long expectedValue = ((Number) fieldSpec.getDefaultNullValue()).longValue();
+ long[] expectedArray = new long[]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ MultiValueResult mvResult = reader.nextLongMV();
+ Assert.assertFalse(mvResult.hasNulls());
+ Assert.assertTrue(Arrays.equals(mvResult.getValues(), expectedArray));
+ }
+
+ // Test random access and array reuse
+ reader.rewind();
+ long[] firstCall = reader.getLongMV(0).getValues();
+ long[] secondCall = reader.getLongMV(1).getValues();
+ Assert.assertTrue(Arrays.equals(firstCall, expectedArray));
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ @Test
+ public void testMultiValueFloatColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.FLOAT, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test sequential access with nextFloatMV
+ float expectedValue = ((Number) fieldSpec.getDefaultNullValue()).floatValue();
+ float[] expectedArray = new float[]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ MultiValueResult mvResult = reader.nextFloatMV();
+ Assert.assertFalse(mvResult.hasNulls());
+ Assert.assertTrue(Arrays.equals(mvResult.getValues(), expectedArray));
+ }
+
+ // Test random access and array reuse
+ reader.rewind();
+ float[] firstCall = reader.getFloatMV(0).getValues();
+ float[] secondCall = reader.getFloatMV(1).getValues();
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ @Test
+ public void testMultiValueDoubleColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.DOUBLE, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test sequential access with nextDoubleMV
+ double expectedValue = ((Number) fieldSpec.getDefaultNullValue()).doubleValue();
+ double[] expectedArray = new double[]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ MultiValueResult mvResult = reader.nextDoubleMV();
+ Assert.assertFalse(mvResult.hasNulls());
+ Assert.assertTrue(Arrays.equals(mvResult.getValues(), expectedArray));
+ }
+
+ // Test random access and array reuse
+ reader.rewind();
+ double[] firstCall = reader.getDoubleMV(0).getValues();
+ double[] secondCall = reader.getDoubleMV(1).getValues();
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ @Test
+ public void testMultiValueStringColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.STRING, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test sequential access with nextStringMV
+ String expectedValue = (String) fieldSpec.getDefaultNullValue();
+ String[] expectedArray = new String[]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ String[] result = reader.nextStringMV();
+ Assert.assertTrue(Arrays.equals(result, expectedArray));
+ }
+
+ // Test random access and array reuse
+ reader.rewind();
+ String[] firstCall = reader.getStringMV(0);
+ String[] secondCall = reader.getStringMV(1);
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ @Test
+ public void testMultiValueBytesColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.BYTES, false);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Test sequential access with nextBytesMV
+ byte[] expectedValue = (byte[]) fieldSpec.getDefaultNullValue();
+ byte[][] expectedArray = new byte[][]{expectedValue};
+ for (int i = 0; i < NUM_DOCS; i++) {
+ byte[][] result = reader.nextBytesMV();
+ Assert.assertEquals(result.length, expectedArray.length);
+ Assert.assertTrue(Arrays.equals(result[0], expectedArray[0]));
+ }
+
+ // Test random access and array reuse
+ reader.rewind();
+ byte[][] firstCall = reader.getBytesMV(0);
+ byte[][] secondCall = reader.getBytesMV(1);
+ Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused");
+
+ reader.close();
+ }
+
+ // ========== Edge Cases and Error Handling ==========
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testNextPastEnd() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 5, fieldSpec);
+
+ // Read all values
+ for (int i = 0; i < 5; i++) {
+ reader.nextInt();
+ }
+
+ // This should throw IllegalStateException
+ reader.nextInt();
+ reader.close();
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testIsNextNullPastEnd() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 5, fieldSpec);
+
+ // Read all values
+ for (int i = 0; i < 5; i++) {
+ reader.nextInt();
+ }
+
+ // This should throw IllegalStateException
+ reader.isNextNull();
+ reader.close();
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testSkipNextPastEnd() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 5, fieldSpec);
+
+ // Read all values
+ for (int i = 0; i < 5; i++) {
+ reader.skipNext();
+ }
+
+ // This should throw IllegalStateException
+ reader.skipNext();
+ reader.close();
+ }
+
+ @Test(expectedExceptions = IndexOutOfBoundsException.class)
+ public void testRandomAccessOutOfBounds() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // This should throw IndexOutOfBoundsException
+ reader.getInt(NUM_DOCS);
+ reader.close();
+ }
+
+ @Test(expectedExceptions = IndexOutOfBoundsException.class)
+ public void testRandomAccessNegativeIndex() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // This should throw IndexOutOfBoundsException
+ reader.getInt(-1);
+ reader.close();
+ }
+
+ @Test
+ public void testIsNullAlwaysReturnsFalse() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, NUM_DOCS, fieldSpec);
+
+ // Default values are never null
+ for (int i = 0; i < NUM_DOCS; i++) {
+ Assert.assertFalse(reader.isNull(i));
+ Assert.assertFalse(reader.isNextNull());
+ reader.skipNext();
+ }
+
+ reader.close();
+ }
+
+ @Test
+ public void testSkipNext() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 10, fieldSpec);
+
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+
+ // Skip every other value
+ for (int i = 0; i < 5; i++) {
+ Assert.assertEquals(reader.nextInt(), expectedValue);
+ reader.skipNext();
+ }
+
+ Assert.assertFalse(reader.hasNext());
+ reader.close();
+ }
+
+ @Test
+ public void testRewind() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 10, fieldSpec);
+
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+
+ // Read all values
+ for (int i = 0; i < 10; i++) {
+ Assert.assertEquals(reader.nextInt(), expectedValue);
+ }
+ Assert.assertFalse(reader.hasNext());
+
+ // Rewind and read again
+ reader.rewind();
+ Assert.assertTrue(reader.hasNext());
+ for (int i = 0; i < 10; i++) {
+ Assert.assertEquals(reader.nextInt(), expectedValue);
+ }
+ Assert.assertFalse(reader.hasNext());
+
+ reader.close();
+ }
+
+ @Test
+ public void testNextObjectMethod() throws IOException {
+ // Test single-value field
+ FieldSpec svFieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader svReader = new DefaultValueColumnReader(COLUMN_NAME, 5, svFieldSpec);
+
+ int expectedSvValue = ((Number) svFieldSpec.getDefaultNullValue()).intValue();
+ for (int i = 0; i < 5; i++) {
+ Object value = svReader.next();
+ Assert.assertEquals(value, expectedSvValue);
+ }
+ svReader.close();
+
+ // Test multi-value field
+ FieldSpec mvFieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, false);
+ DefaultValueColumnReader mvReader = new DefaultValueColumnReader(COLUMN_NAME, 5, mvFieldSpec);
+
+ int expectedMvValue = ((Number) mvFieldSpec.getDefaultNullValue()).intValue();
+ Object[] expectedArray = new Object[]{expectedMvValue};
+ for (int i = 0; i < 5; i++) {
+ Object value = mvReader.next();
+ Assert.assertTrue(value instanceof Object[]);
+ Assert.assertTrue(Arrays.equals((Object[]) value, expectedArray));
+ }
+ mvReader.close();
+ }
+
+ @Test
+ public void testGetTotalDocs() throws IOException {
+ for (int numDocs : new int[]{0, 1, 10, 100, 1000}) {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, numDocs, fieldSpec);
+ Assert.assertEquals(reader.getTotalDocs(), numDocs);
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testEmptyColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 0, fieldSpec);
+
+ Assert.assertEquals(reader.getTotalDocs(), 0);
+ Assert.assertFalse(reader.hasNext());
+
+ reader.close();
+ }
+
+ @Test
+ public void testSingleDocColumn() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 1, fieldSpec);
+
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+
+ Assert.assertTrue(reader.hasNext());
+ Assert.assertEquals(reader.nextInt(), expectedValue);
+ Assert.assertFalse(reader.hasNext());
+
+ reader.rewind();
+ Assert.assertEquals(reader.getInt(0), expectedValue);
+
+ reader.close();
+ }
+
+ @Test
+ public void testAllTypeIndicators() throws IOException {
+ // Test INT
+ FieldSpec intSpec = new DimensionFieldSpec("int_col", FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader intReader = new DefaultValueColumnReader("int_col", 1, intSpec);
+ Assert.assertTrue(intReader.isInt());
+ Assert.assertFalse(intReader.isLong());
+ Assert.assertFalse(intReader.isFloat());
+ Assert.assertFalse(intReader.isDouble());
+ Assert.assertFalse(intReader.isString());
+ Assert.assertFalse(intReader.isBytes());
+ intReader.close();
+
+ // Test LONG
+ FieldSpec longSpec = new DimensionFieldSpec("long_col", FieldSpec.DataType.LONG, true);
+ DefaultValueColumnReader longReader = new DefaultValueColumnReader("long_col", 1, longSpec);
+ Assert.assertFalse(longReader.isInt());
+ Assert.assertTrue(longReader.isLong());
+ longReader.close();
+
+ // Test FLOAT
+ FieldSpec floatSpec = new DimensionFieldSpec("float_col", FieldSpec.DataType.FLOAT, true);
+ DefaultValueColumnReader floatReader = new DefaultValueColumnReader("float_col", 1, floatSpec);
+ Assert.assertTrue(floatReader.isFloat());
+ Assert.assertFalse(floatReader.isDouble());
+ floatReader.close();
+
+ // Test DOUBLE
+ FieldSpec doubleSpec = new DimensionFieldSpec("double_col", FieldSpec.DataType.DOUBLE, true);
+ DefaultValueColumnReader doubleReader = new DefaultValueColumnReader("double_col", 1, doubleSpec);
+ Assert.assertFalse(doubleReader.isFloat());
+ Assert.assertTrue(doubleReader.isDouble());
+ doubleReader.close();
+
+ // Test STRING
+ FieldSpec stringSpec = new DimensionFieldSpec("string_col", FieldSpec.DataType.STRING, true);
+ DefaultValueColumnReader stringReader = new DefaultValueColumnReader("string_col", 1, stringSpec);
+ Assert.assertTrue(stringReader.isString());
+ Assert.assertFalse(stringReader.isBytes());
+ stringReader.close();
+
+ // Test BYTES
+ FieldSpec bytesSpec = new DimensionFieldSpec("bytes_col", FieldSpec.DataType.BYTES, true);
+ DefaultValueColumnReader bytesReader = new DefaultValueColumnReader("bytes_col", 1, bytesSpec);
+ Assert.assertFalse(bytesReader.isString());
+ Assert.assertTrue(bytesReader.isBytes());
+ bytesReader.close();
+ }
+
+ @Test
+ public void testMultipleReadsFromSamePosition() throws IOException {
+ FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN_NAME, FieldSpec.DataType.INT, true);
+ DefaultValueColumnReader reader = new DefaultValueColumnReader(COLUMN_NAME, 10, fieldSpec);
+
+ int expectedValue = ((Number) fieldSpec.getDefaultNullValue()).intValue();
+
+ // Read from the same random access position multiple times
+ for (int i = 0; i < 5; i++) {
+ Assert.assertEquals(reader.getInt(0), expectedValue);
+ Assert.assertEquals(reader.getInt(5), expectedValue);
+ Assert.assertEquals(reader.getInt(9), expectedValue);
+ }
+
+ reader.close();
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/InMemoryColumnReader.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/InMemoryColumnReader.java
new file mode 100644
index 0000000000..adda7e5228
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/InMemoryColumnReader.java
@@ -0,0 +1,258 @@
+/**
+ * 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 javax.annotation.Nullable;
+import org.apache.pinot.spi.data.readers.ColumnReader;
+import org.apache.pinot.spi.data.readers.MultiValueResult;
+
+
+/**
+ * Mock ColumnReader for testing with configurable data.
+ */
+public class InMemoryColumnReader implements ColumnReader {
+ private final String _columnName;
+ private final Object[] _values;
+ private final boolean[] _isNull;
+ private final boolean _isSingleValue;
+ private final Class _dataType;
+ private int _currentIndex = 0;
+
+ public InMemoryColumnReader(String columnName, Object[] values, boolean[] isNull, boolean isSingleValue,
+ Class dataType) {
+ _columnName = columnName;
+ _values = values;
+ _isNull = isNull;
+ _isSingleValue = isSingleValue;
+ _dataType = dataType;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return _currentIndex < _values.length;
+ }
+
+ @Nullable
+ @Override
+ public Object next() {
+ if (!hasNext()) {
+ throw new IllegalStateException("No more values");
+ }
+ Object value = _values[_currentIndex];
+ _currentIndex++;
+ return value;
+ }
+
+ @Override
+ public boolean isNextNull() {
+ return hasNext() && _isNull[_currentIndex];
+ }
+
+ @Override
+ public void skipNext() {
+ if (hasNext()) {
+ _currentIndex++;
+ }
+ }
+
+ @Override
+ public boolean isSingleValue() {
+ return _isSingleValue;
+ }
+
+ @Override
+ public boolean isInt() {
+ return _dataType == Integer.class;
+ }
+
+ @Override
+ public boolean isLong() {
+ return _dataType == Long.class;
+ }
+
+ @Override
+ public boolean isFloat() {
+ return _dataType == Float.class;
+ }
+
+ @Override
+ public boolean isDouble() {
+ return _dataType == Double.class;
+ }
+
+ @Override
+ public boolean isString() {
+ return _dataType == String.class;
+ }
+
+ @Override
+ public boolean isBytes() {
+ return _dataType == byte[].class;
+ }
+
+ @Override
+ public int nextInt() {
+ return (Integer) next();
+ }
+
+ @Override
+ public long nextLong() {
+ return (Long) next();
+ }
+
+ @Override
+ public float nextFloat() {
+ return (Float) next();
+ }
+
+ @Override
+ public double nextDouble() {
+ return (Double) next();
+ }
+
+ @Override
+ public String nextString() {
+ return (String) next();
+ }
+
+ @Override
+ public byte[] nextBytes() {
+ return (byte[]) next();
+ }
+
+ @Override
+ public MultiValueResult nextIntMV() {
+ return MultiValueResult.of((int[]) next(), null);
+ }
+
+ @Override
+ public MultiValueResult nextLongMV() {
+ return MultiValueResult.of((long[]) next(), null);
+ }
+
+ @Override
+ public MultiValueResult nextFloatMV() {
+ return MultiValueResult.of((float[]) next(), null);
+ }
+
+ @Override
+ public MultiValueResult nextDoubleMV() {
+ return MultiValueResult.of((double[]) next(), null);
+ }
+
+ @Override
+ public String[] nextStringMV() {
+ return (String[]) next();
+ }
+
+ @Override
+ public byte[][] nextBytesMV() {
+ return (byte[][]) next();
+ }
+
+ @Override
+ public void rewind() {
+ _currentIndex = 0;
+ }
+
+ @Override
+ public String getColumnName() {
+ return _columnName;
+ }
+
+ @Override
+ public int getTotalDocs() {
+ return _values.length;
+ }
+
+ @Override
+ public boolean isNull(int docId) {
+ return _isNull[docId];
+ }
+
+ @Override
+ public int getInt(int docId) {
+ return (Integer) _values[docId];
+ }
+
+ @Override
+ public long getLong(int docId) {
+ return (Long) _values[docId];
+ }
+
+ @Override
+ public float getFloat(int docId) {
+ return (Float) _values[docId];
+ }
+
+ @Override
+ public double getDouble(int docId) {
+ return (Double) _values[docId];
+ }
+
+ @Override
+ public String getString(int docId) {
+ return (String) _values[docId];
+ }
+
+ @Override
+ public byte[] getBytes(int docId) {
+ return (byte[]) _values[docId];
+ }
+
+ @Override
+ public Object getValue(int docId) {
+ return _values[docId];
+ }
+
+ @Override
+ public MultiValueResult getIntMV(int docId) {
+ return MultiValueResult.of((int[]) _values[docId], null);
+ }
+
+ @Override
+ public MultiValueResult getLongMV(int docId) {
+ return MultiValueResult.of((long[]) _values[docId], null);
+ }
+
+ @Override
+ public MultiValueResult getFloatMV(int docId) {
+ return MultiValueResult.of((float[]) _values[docId], null);
+ }
+
+ @Override
+ public MultiValueResult getDoubleMV(int docId) {
+ return MultiValueResult.of((double[]) _values[docId], null);
+ }
+
+ @Override
+ public String[] getStringMV(int docId) {
+ return (String[]) _values[docId];
+ }
+
+ @Override
+ public byte[][] getBytesMV(int docId) {
+ return (byte[][]) _values[docId];
+ }
+
+ @Override
+ public void close() {
+ // No-op
+ }
+}
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java
new file mode 100644
index 0000000000..f3e8d0141f
--- /dev/null
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java
@@ -0,0 +1,963 @@
+/**
+ * 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.io.File;
+import java.io.IOException;
+import java.util.function.Function;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.ColumnarSegmentBuildingTestBase;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.MultiValueResult;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+
+/**
+ * Comprehensive tests for PinotSegmentColumnReaderImpl random access methods.
+ *
+ * This test validates:
+ *
+ * - Single-value accessor methods (getInt, getLong, getFloat, getDouble, getString, getBytes)
+ * - Multi-value accessor methods (getIntMV, getLongMV, getFloatMV, getDoubleMV, getStringMV, getBytesMV)
+ * - Metadata methods (getTotalDocs, isNull)
+ * - Boundary conditions and exception handling
+ * - Consistency between iterator and random access patterns
+ * - Multiple readers for the same column
+ * - Both dictionary-encoded and raw forward index columns
+ *
+ */
+public class PinotSegmentColumnReaderImplTest extends ColumnarSegmentBuildingTestBase {
+
+ /**
+ * Enum representing different segment index configurations for testing.
+ */
+ enum IndexType {
+ DICT_ENCODED,
+ RAW_INDEX,
+ NULLABLE
+ }
+
+ /**
+ * Functional interface for single-value getters that can throw IOException.
+ */
+ @FunctionalInterface
+ interface SingleValueGetter {
+ Object get(PinotSegmentColumnReaderImpl reader, int docId)
+ throws IOException;
+ }
+
+ /**
+ * Functional interface for single-value sequential getters that can throw IOException.
+ */
+ @FunctionalInterface
+ interface SingleValueSequentialGetter {
+ Object get(PinotSegmentColumnReaderImpl reader)
+ throws IOException;
+ }
+
+ /**
+ * Functional interface for multi-value getters that can throw IOException.
+ */
+ @FunctionalInterface
+ interface MultiValueGetter {
+ Object get(PinotSegmentColumnReaderImpl reader, int docId)
+ throws IOException;
+ }
+
+ /**
+ * Functional interface for multi-value sequential getters that can throw IOException.
+ */
+ @FunctionalInterface
+ interface MultiValueSequentialGetter {
+ Object get(PinotSegmentColumnReaderImpl reader)
+ throws IOException;
+ }
+
+ private ImmutableSegment _dictEncodedSegment;
+ private ImmutableSegment _rawIndexSegment;
+ private ImmutableSegment _nullableSegment;
+
+ @BeforeClass
+ @Override
+ public void setUp()
+ throws IOException {
+ super.setUp();
+ // Create test segments with different configurations
+ try {
+ // Dictionary-encoded segment (default)
+ File dictEncodedSegmentDir = createRowMajorSegment();
+ _dictEncodedSegment = ImmutableSegmentLoader.load(dictEncodedSegmentDir, ReadMode.mmap);
+
+ // Raw forward index segment (no dictionary)
+ TableConfig noDictConfig = createTableConfigWithNoDictionary();
+ File rawIndexSegmentDir = createRowMajorSegmentWithConfig(noDictConfig, _originalSchema, "rawIndexSegment");
+ _rawIndexSegment = ImmutableSegmentLoader.load(rawIndexSegmentDir, ReadMode.mmap);
+
+ // Nullable segment (allows null values)
+ Schema nullableSchema = createNullableSchema();
+ File nullableSegmentDir = createRowMajorSegmentWithConfig(_tableConfig, nullableSchema, "nullableSegment");
+ _nullableSegment = ImmutableSegmentLoader.load(nullableSegmentDir, ReadMode.mmap);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @AfterClass
+ public void tearDownTest()
+ throws Exception {
+ if (_dictEncodedSegment != null) {
+ _dictEncodedSegment.destroy();
+ }
+ if (_rawIndexSegment != null) {
+ _rawIndexSegment.destroy();
+ }
+ if (_nullableSegment != null) {
+ _nullableSegment.destroy();
+ }
+ super.tearDown();
+ }
+
+ @Test
+ public void testGetTotalDocs()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ Assert.assertEquals(reader.getTotalDocs(), _testData.size());
+ Assert.assertEquals(reader.getTotalDocs(), _dictEncodedSegment.getSegmentMetadata().getTotalDocs());
+
+ reader.close();
+ }
+
+ /**
+ * Helper method to get the appropriate segment and schema based on IndexType.
+ */
+ private ImmutableSegment getSegment(IndexType indexType) {
+ switch (indexType) {
+ case DICT_ENCODED:
+ return _dictEncodedSegment;
+ case RAW_INDEX:
+ return _rawIndexSegment;
+ case NULLABLE:
+ return _nullableSegment;
+ default:
+ throw new IllegalArgumentException("Unknown index type: " + indexType);
+ }
+ }
+
+ /**
+ * Helper method to get the appropriate schema based on IndexType.
+ */
+ private Schema getSchema(IndexType indexType) {
+ switch (indexType) {
+ case DICT_ENCODED:
+ case RAW_INDEX:
+ return _originalSchema;
+ case NULLABLE:
+ return createNullableSchema();
+ default:
+ throw new IllegalArgumentException("Unknown index type: " + indexType);
+ }
+ }
+
+ /**
+ * Data provider for single-value accessor methods.
+ * @return test parameters: column name, random access getter, sequential getter,
+ * value converter function (required for float), index type
+ */
+ @DataProvider(name = "singleValueAccessorProvider")
+ public Object[][] singleValueAccessorProvider() {
+ // Base column configurations
+ Object[][] baseConfigs = new Object[][] {
+ {INT_COL_1, (SingleValueGetter) PinotSegmentColumnReaderImpl::getInt,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextInt, null},
+ {LONG_COL, (SingleValueGetter) PinotSegmentColumnReaderImpl::getLong,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextLong, null},
+ {FLOAT_COL, (SingleValueGetter) PinotSegmentColumnReaderImpl::getFloat,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextFloat,
+ (Function) val -> val instanceof Double ? ((Double) val).floatValue() : val},
+ {DOUBLE_COL, (SingleValueGetter) PinotSegmentColumnReaderImpl::getDouble,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextDouble, null},
+ {STRING_COL_1, (SingleValueGetter) PinotSegmentColumnReaderImpl::getString,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextString, null},
+ {BYTES_COL, (SingleValueGetter) PinotSegmentColumnReaderImpl::getBytes,
+ (SingleValueSequentialGetter) PinotSegmentColumnReaderImpl::nextBytes, null},
+ };
+
+ // Create test cases for all index types
+ IndexType[] indexTypes = IndexType.values();
+ Object[][] result = new Object[baseConfigs.length * indexTypes.length][];
+ int idx = 0;
+ for (Object[] baseConfig : baseConfigs) {
+ for (IndexType indexType : indexTypes) {
+ result[idx++] = new Object[]{baseConfig[0], baseConfig[1], baseConfig[2], baseConfig[3], indexType};
+ }
+ }
+ return result;
+ }
+
+ @Test(dataProvider = "singleValueAccessorProvider")
+ public void testSingleValueAccessors(String columnName, SingleValueGetter randomAccessGetter,
+ SingleValueSequentialGetter sequentialGetter, Function valueConverter, IndexType indexType)
+ throws Exception {
+ ImmutableSegment segment = getSegment(indexType);
+ Schema schema = getSchema(indexType);
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(segment, columnName);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+ Assert.assertEquals(totalDocs, _testData.size());
+
+ // Get field spec and data type for the column
+ FieldSpec fieldSpec = schema.getFieldSpecFor(columnName);
+ Object defaultValue = fieldSpec.getDefaultNullValue();
+ FieldSpec.DataType dataType = fieldSpec.getDataType();
+
+ // Test type indicator methods
+ Assert.assertEquals(reader.isInt(), dataType == FieldSpec.DataType.INT,
+ "isInt() should return " + (dataType == FieldSpec.DataType.INT) + " for " + columnName);
+ Assert.assertEquals(reader.isLong(), dataType == FieldSpec.DataType.LONG,
+ "isLong() should return " + (dataType == FieldSpec.DataType.LONG) + " for " + columnName);
+ Assert.assertEquals(reader.isFloat(), dataType == FieldSpec.DataType.FLOAT,
+ "isFloat() should return " + (dataType == FieldSpec.DataType.FLOAT) + " for " + columnName);
+ Assert.assertEquals(reader.isDouble(), dataType == FieldSpec.DataType.DOUBLE,
+ "isDouble() should return " + (dataType == FieldSpec.DataType.DOUBLE) + " for " + columnName);
+ Assert.assertEquals(reader.isString(), dataType == FieldSpec.DataType.STRING,
+ "isString() should return " + (dataType == FieldSpec.DataType.STRING) + " for " + columnName);
+ Assert.assertEquals(reader.isBytes(), dataType == FieldSpec.DataType.BYTES,
+ "isBytes() should return " + (dataType == FieldSpec.DataType.BYTES) + " for " + columnName);
+
+ // Test reading all documents using random access (Pattern 3)
+ for (int docId = 0; docId < totalDocs; docId++) {
+ GenericRow expectedRow = _testData.get(docId);
+ Object expectedValue = expectedRow.getValue(columnName);
+
+ // For nullable segments, check if the value is actually null
+ if (indexType == IndexType.NULLABLE && expectedValue == null) {
+ Assert.assertTrue(reader.isNull(docId),
+ "isNull() should return true for null value at docId " + docId);
+ continue;
+ }
+
+ Object actualValue = randomAccessGetter.get(reader, docId);
+
+ if (expectedValue == null) {
+ // Null values are replaced with default during segment creation (non-nullable)
+ if (actualValue instanceof byte[]) {
+ Assert.assertEquals(actualValue, defaultValue, "Null should be replaced with default at docId " + docId);
+ } else if (actualValue instanceof Double) {
+ Assert.assertEquals((Double) actualValue, (Double) defaultValue, 0.0001,
+ "Null should be replaced with default at docId " + docId);
+ } else if (actualValue instanceof Float) {
+ Assert.assertEquals((Float) actualValue, (Float) defaultValue, 0.0001f,
+ "Null should be replaced with default at docId " + docId);
+ } else {
+ Assert.assertEquals(actualValue, defaultValue, "Null should be replaced with default at docId " + docId);
+ }
+ } else {
+ // Convert expected value if converter is provided
+ Object convertedExpectedValue = valueConverter != null ? valueConverter.apply(expectedValue) : expectedValue;
+ if (actualValue instanceof Double) {
+ Assert.assertEquals((Double) actualValue, (Double) convertedExpectedValue, 0.0001,
+ "Value mismatch at docId " + docId);
+ } else if (actualValue instanceof Float) {
+ Assert.assertEquals((Float) actualValue, (Float) convertedExpectedValue, 0.0001f,
+ "Value mismatch at docId " + docId);
+ } else {
+ Assert.assertEquals(actualValue, convertedExpectedValue, "Value mismatch at docId " + docId);
+ }
+ }
+ }
+
+ // Test sequential iteration with type-specific methods (Pattern 2) - only for columns with sequential getters
+ if (sequentialGetter != null) {
+ reader.rewind();
+ int docId = 0;
+ while (reader.hasNext()) {
+ GenericRow expectedRow = _testData.get(docId);
+ Object expectedValue = expectedRow.getValue(columnName);
+
+ if (indexType == IndexType.NULLABLE && expectedValue == null) {
+ Assert.assertTrue(reader.isNextNull(),
+ "isNextNull() should return true for null value at docId " + docId);
+ reader.skipNext();
+ } else if (reader.isNextNull()) {
+ reader.skipNext();
+ // For non-nullable segments, null is replaced with default
+ Assert.assertNull(expectedValue, "Expected null value at docId " + docId);
+ } else {
+ Object actualValue = sequentialGetter.get(reader);
+
+ if (expectedValue == null) {
+ // Null values are replaced with default during segment creation (non-nullable)
+ if (actualValue instanceof byte[]) {
+ Assert.assertEquals(actualValue, defaultValue,
+ "Null should be replaced with default at docId " + docId + " (sequential)");
+ } else if (actualValue instanceof Double) {
+ Assert.assertEquals((Double) actualValue, (Double) defaultValue, 0.0001,
+ "Null should be replaced with default at docId " + docId + " (sequential)");
+ } else if (actualValue instanceof Float) {
+ Assert.assertEquals((Float) actualValue, (Float) defaultValue, 0.0001f,
+ "Null should be replaced with default at docId " + docId + " (sequential)");
+ } else {
+ Assert.assertEquals(actualValue, defaultValue,
+ "Null should be replaced with default at docId " + docId + " (sequential)");
+ }
+ } else {
+ // Convert expected value if converter is provided
+ Object convertedValue = valueConverter != null ? valueConverter.apply(expectedValue) : expectedValue;
+ if (actualValue instanceof Double) {
+ Assert.assertEquals((Double) actualValue, (Double) convertedValue, 0.0001,
+ "Value mismatch at docId " + docId + " (sequential)");
+ } else if (actualValue instanceof Float) {
+ Assert.assertEquals((Float) actualValue, (Float) convertedValue, 0.0001f,
+ "Value mismatch at docId " + docId + " (sequential)");
+ } else {
+ Assert.assertEquals(actualValue, convertedValue,
+ "Value mismatch at docId " + docId + " (sequential)");
+ }
+ }
+ }
+ docId++;
+ }
+ Assert.assertEquals(docId, totalDocs, "Should have iterated through all documents");
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+
+ /**
+ * Data provider for multi-value accessor methods.
+ * @return test parameters: column name, random access getter function,
+ * sequential getter function, array converter function, index type
+ */
+ @DataProvider(name = "multiValueAccessorProvider")
+ public Object[][] multiValueAccessorProvider() {
+ // Base column configurations
+ // For primitive MV types, we need to extract .getValues() from MultiValueResult
+ // and verify that hasNulls() is false (nulls are removed by NullValueTransformer for MV primitive types)
+ Object[][] baseConfigs = new Object[][] {
+ {MV_INT_COL, (MultiValueGetter) (reader, docId) -> {
+ MultiValueResult result = reader.getIntMV(docId);
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ },
+ (MultiValueSequentialGetter) reader -> {
+ MultiValueResult result = reader.nextIntMV();
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ }, null},
+ {MV_LONG_COL, (MultiValueGetter) (reader, docId) -> {
+ MultiValueResult result = reader.getLongMV(docId);
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ },
+ (MultiValueSequentialGetter) reader -> {
+ MultiValueResult result = reader.nextLongMV();
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ }, null},
+ {MV_FLOAT_COL, (MultiValueGetter) (reader, docId) -> {
+ MultiValueResult result = reader.getFloatMV(docId);
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ },
+ (MultiValueSequentialGetter) reader -> {
+ MultiValueResult result = reader.nextFloatMV();
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ },
+ (Function) expectedArray -> {
+ float[] expectedFloatArray = new float[expectedArray.length];
+ for (int i = 0; i < expectedArray.length; i++) {
+ expectedFloatArray[i] = (Float) expectedArray[i];
+ }
+ return expectedFloatArray;
+ } },
+ {MV_DOUBLE_COL, (MultiValueGetter) (reader, docId) -> {
+ MultiValueResult result = reader.getDoubleMV(docId);
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ },
+ (MultiValueSequentialGetter) reader -> {
+ MultiValueResult result = reader.nextDoubleMV();
+ Assert.assertFalse(result.hasNulls(), "Multi-value primitive types should not have nulls");
+ return result.getValues();
+ }, null},
+ {MV_STRING_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getStringMV,
+ (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextStringMV, null},
+ {MV_BYTES_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getBytesMV,
+ (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextBytesMV, null},
+ };
+
+ // Create test cases for all index types
+ IndexType[] indexTypes = IndexType.values();
+ Object[][] result = new Object[baseConfigs.length * indexTypes.length][];
+ int idx = 0;
+ for (Object[] baseConfig : baseConfigs) {
+ for (IndexType indexType : indexTypes) {
+ result[idx++] = new Object[]{baseConfig[0], baseConfig[1], baseConfig[2], baseConfig[3], indexType};
+ }
+ }
+ return result;
+ }
+
+ @Test(dataProvider = "multiValueAccessorProvider")
+ public void testMultiValueAccessors(String columnName, MultiValueGetter randomAccessGetter,
+ MultiValueSequentialGetter sequentialGetter, Function arrayConverter, IndexType indexType)
+ throws Exception {
+ ImmutableSegment segment = getSegment(indexType);
+ Schema schema = getSchema(indexType);
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(segment, columnName);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // Get default value for MV type (wrapped in array)
+ FieldSpec fieldSpec = schema.getFieldSpecFor(columnName);
+ Object defaultNullValue = fieldSpec.getDefaultNullValue();
+ FieldSpec.DataType dataType = fieldSpec.getDataType();
+ Object defaultValue;
+
+ // Test type indicator methods
+ Assert.assertEquals(reader.isInt(), dataType == FieldSpec.DataType.INT,
+ "isInt() should return " + (dataType == FieldSpec.DataType.INT) + " for " + columnName);
+ Assert.assertEquals(reader.isLong(), dataType == FieldSpec.DataType.LONG,
+ "isLong() should return " + (dataType == FieldSpec.DataType.LONG) + " for " + columnName);
+ Assert.assertEquals(reader.isFloat(), dataType == FieldSpec.DataType.FLOAT,
+ "isFloat() should return " + (dataType == FieldSpec.DataType.FLOAT) + " for " + columnName);
+ Assert.assertEquals(reader.isDouble(), dataType == FieldSpec.DataType.DOUBLE,
+ "isDouble() should return " + (dataType == FieldSpec.DataType.DOUBLE) + " for " + columnName);
+ Assert.assertEquals(reader.isString(), dataType == FieldSpec.DataType.STRING,
+ "isString() should return " + (dataType == FieldSpec.DataType.STRING) + " for " + columnName);
+ Assert.assertEquals(reader.isBytes(), dataType == FieldSpec.DataType.BYTES,
+ "isBytes() should return " + (dataType == FieldSpec.DataType.BYTES) + " for " + columnName);
+
+ // Create default array based on type
+ switch (fieldSpec.getDataType()) {
+ case INT:
+ defaultValue = new int[]{((Number) defaultNullValue).intValue()};
+ break;
+ case LONG:
+ defaultValue = new long[]{((Number) defaultNullValue).longValue()};
+ break;
+ case FLOAT:
+ defaultValue = new float[]{((Number) defaultNullValue).floatValue()};
+ break;
+ case DOUBLE:
+ defaultValue = new double[]{((Number) defaultNullValue).doubleValue()};
+ break;
+ case STRING:
+ defaultValue = new String[]{(String) defaultNullValue};
+ break;
+ case BYTES:
+ defaultValue = new byte[][]{(byte[]) defaultNullValue};
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported data type: " + fieldSpec.getDataType());
+ }
+
+ // Test reading all documents using random access
+ for (int docId = 0; docId < totalDocs; docId++) {
+ GenericRow expectedRow = _testData.get(docId);
+ Object expectedValue = expectedRow.getValue(columnName);
+
+ // For nullable segments, check if the value is actually null
+ if (indexType == IndexType.NULLABLE && expectedValue == null) {
+ Assert.assertTrue(reader.isNull(docId),
+ "isNull() should return true for null value at docId " + docId);
+ continue;
+ }
+
+ Object actualValue = randomAccessGetter.get(reader, docId);
+
+ if (expectedValue == null) {
+ // Null values are replaced with default array during segment creation (non-nullable)
+ Assert.assertEquals(actualValue, defaultValue,
+ "Null should be replaced with default array at docId " + docId);
+ } else {
+ // Convert expected Object[] to appropriate type array for comparison
+ Object[] expectedArray = (Object[]) expectedValue;
+ Object expectedConvertedArray = arrayConverter != null ? arrayConverter.apply(expectedArray) : expectedArray;
+
+ Assert.assertEquals(actualValue, expectedConvertedArray,
+ "MV array mismatch at docId " + docId);
+ }
+ }
+
+ // Test sequential iteration with type-specific MV methods
+ reader.rewind();
+ int docId = 0;
+ while (reader.hasNext()) {
+ GenericRow expectedRow = _testData.get(docId);
+ Object expectedValue = expectedRow.getValue(columnName);
+
+ if (indexType == IndexType.NULLABLE && expectedValue == null) {
+ Assert.assertTrue(reader.isNextNull(),
+ "isNextNull() should return true for null value at docId " + docId);
+ reader.skipNext();
+ } else if (reader.isNextNull()) {
+ reader.skipNext();
+ // For non-nullable segments, null is replaced with default
+ Assert.assertNull(expectedValue, "Expected null value at docId " + docId);
+ } else {
+ Object actualValue = sequentialGetter.get(reader);
+
+ if (expectedValue == null) {
+ // Null values are replaced with default array during segment creation (non-nullable)
+ Assert.assertEquals(actualValue, defaultValue,
+ "Null should be replaced with default array at docId " + docId + " (sequential)");
+ } else {
+ // Convert expected Object[] to appropriate type array for comparison
+ Object[] expectedArray = (Object[]) expectedValue;
+ Object convertedArray = arrayConverter != null ? arrayConverter.apply(expectedArray) : expectedArray;
+
+ Assert.assertEquals(actualValue, convertedArray,
+ "MV array mismatch at docId " + docId + " (sequential)");
+ }
+ }
+ docId++;
+ }
+
+ Assert.assertEquals(docId, totalDocs, "Should have iterated through all documents");
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testIsNull()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, STRING_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // In Pinot segments, nulls are replaced with defaults, so isNull should always return false
+ for (int docId = 0; docId < totalDocs; docId++) {
+ boolean actualNull = reader.isNull(docId);
+ Assert.assertFalse(actualNull, "isNull() should return false (nulls replaced with defaults) at docId " + docId);
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testBoundaryConditions()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+ Assert.assertTrue(totalDocs > 0, "Expected non-empty segment");
+
+ // Test first document (docId = 0)
+ int firstValue = reader.getInt(0);
+ Object expectedFirst = _testData.get(0).getValue(INT_COL_1);
+ if (expectedFirst != null) {
+ Assert.assertEquals(firstValue, expectedFirst, "First document value mismatch");
+ }
+
+ // Test last document (docId = totalDocs - 1)
+ int lastDocId = totalDocs - 1;
+ int lastValue = reader.getInt(lastDocId);
+ Object expectedLast = _testData.get(lastDocId).getValue(INT_COL_1);
+ if (expectedLast != null) {
+ Assert.assertEquals(lastValue, expectedLast, "Last document value mismatch");
+ }
+
+ // Test docId = -1 (should throw IndexOutOfBoundsException)
+ try {
+ reader.getInt(-1);
+ Assert.fail("Expected IndexOutOfBoundsException for docId = -1");
+ } catch (IndexOutOfBoundsException e) {
+ // Expected
+ }
+
+ // Test docId = totalDocs (should throw IndexOutOfBoundsException)
+ try {
+ reader.getInt(totalDocs);
+ Assert.fail("Expected IndexOutOfBoundsException for docId = totalDocs");
+ } catch (IndexOutOfBoundsException e) {
+ // Expected
+ }
+
+ // Test docId = totalDocs + 100 (should throw IndexOutOfBoundsException)
+ try {
+ reader.getInt(totalDocs + 100);
+ Assert.fail("Expected IndexOutOfBoundsException for docId > totalDocs");
+ } catch (IndexOutOfBoundsException e) {
+ // Expected
+ }
+
+ // Test isNull with out of bounds docId
+ try {
+ reader.isNull(-1);
+ Assert.fail("Expected IndexOutOfBoundsException for isNull(-1)");
+ } catch (IndexOutOfBoundsException e) {
+ // Expected
+ }
+
+ try {
+ reader.isNull(totalDocs);
+ Assert.fail("Expected IndexOutOfBoundsException for isNull(totalDocs)");
+ } catch (IndexOutOfBoundsException e) {
+ // Expected
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testIteratorAndRandomAccessConsistency()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // First, read all values using random access
+ Integer[] randomAccessValues = new Integer[totalDocs];
+ for (int docId = 0; docId < totalDocs; docId++) {
+ randomAccessValues[docId] = reader.getInt(docId);
+ }
+
+ // Now read using iterator pattern
+ reader.rewind();
+ Integer[] iteratorValues = new Integer[totalDocs];
+ int index = 0;
+ while (reader.hasNext()) {
+ Object value = reader.next();
+ iteratorValues[index++] = (Integer) value;
+ }
+
+ Assert.assertEquals(index, totalDocs, "Iterator should read all documents");
+
+ // Compare values
+ for (int i = 0; i < totalDocs; i++) {
+ Assert.assertEquals(iteratorValues[i], randomAccessValues[i],
+ "Iterator and random access values differ at index " + i);
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testIteratorAndRandomAccessConsistencyMV()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, MV_INT_COL);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // First, read all values using random access
+ int[][] randomAccessValues = new int[totalDocs][];
+ for (int docId = 0; docId < totalDocs; docId++) {
+ MultiValueResult result = reader.getIntMV(docId);
+ randomAccessValues[docId] = result.getValues();
+ }
+
+ // Now read using iterator pattern
+ reader.rewind();
+ Object[] iteratorValues = new Object[totalDocs];
+ int index = 0;
+ while (reader.hasNext()) {
+ Object value = reader.next();
+ iteratorValues[index++] = value;
+ }
+
+ Assert.assertEquals(index, totalDocs, "Iterator should read all documents");
+
+ // Compare values
+ for (int i = 0; i < totalDocs; i++) {
+ // Convert iterator value (Integer[]) to int[]
+ Integer[] iteratorArray = (Integer[]) iteratorValues[i];
+ int[] iteratorIntArray = new int[iteratorArray.length];
+ for (int j = 0; j < iteratorArray.length; j++) {
+ iteratorIntArray[j] = iteratorArray[j];
+ }
+
+ Assert.assertEquals(iteratorIntArray, randomAccessValues[i],
+ "Iterator and random access MV values differ at index " + i);
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testMultipleReadersForSameColumn()
+ throws Exception {
+ // Test that multiple readers can be created for the same column and read independently
+ PinotSegmentColumnReaderImpl reader1 = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+ PinotSegmentColumnReaderImpl reader2 = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader1.getTotalDocs();
+ Assert.assertEquals(reader2.getTotalDocs(), totalDocs);
+
+ // Read from both readers at different positions
+ for (int docId = 0; docId < Math.min(10, totalDocs); docId++) {
+ int value1 = reader1.getInt(docId);
+ int value2 = reader2.getInt(docId);
+ Assert.assertEquals(value1, value2, "Values should match at docId " + docId);
+ }
+ } finally {
+ reader1.close();
+ reader2.close();
+ }
+ }
+
+ @Test
+ public void testColumnName()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, STRING_COL_1);
+
+ try {
+ Assert.assertEquals(reader.getColumnName(), STRING_COL_1);
+ } finally {
+ reader.close();
+ }
+ }
+
+ /**
+ * Test isNextNull() method behavior.
+ */
+ @Test
+ public void testIsNextNull()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+ int docId = 0;
+
+ // For non-nullable segments, isNextNull should always return false
+ while (reader.hasNext()) {
+ boolean isNull = reader.isNextNull();
+ boolean expectedNull = reader.isNull(docId);
+ Assert.assertEquals(isNull, expectedNull, "isNextNull() mismatch at docId " + docId);
+ // Verify isNextNull doesn't advance the iterator
+ if (!isNull) {
+ int value = reader.nextInt();
+ Assert.assertEquals(value, reader.getInt(docId));
+ } else {
+ reader.skipNext();
+ }
+ docId++;
+ }
+
+ Assert.assertEquals(docId, totalDocs);
+ } finally {
+ reader.close();
+ }
+ }
+
+ /**
+ * Test skipNext() method behavior.
+ */
+ @Test
+ public void testSkipNext()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+ // Skip every other value
+ int docId = 0;
+ int valuesRead = 0;
+ while (reader.hasNext()) {
+ if (docId % 2 == 0) {
+ // Read even-indexed documents
+ int value = reader.nextInt();
+ Assert.assertEquals(value, reader.getInt(docId));
+ valuesRead++;
+ } else {
+ // Skip odd-indexed documents
+ reader.skipNext();
+ }
+ docId++;
+ }
+
+ Assert.assertEquals(docId, totalDocs);
+ Assert.assertEquals(valuesRead, (totalDocs + 1) / 2, "Should have read half the documents");
+ } finally {
+ reader.close();
+ }
+ }
+
+ /**
+ * Test that calling next methods after hasNext returns false throws exception.
+ */
+ @Test
+ public void testNextMethodsAfterEnd()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ // Read all values
+ while (reader.hasNext()) {
+ reader.nextInt();
+ }
+
+ // Now try to read more - should throw IllegalStateException
+ Assert.assertFalse(reader.hasNext());
+
+ try {
+ reader.nextInt();
+ Assert.fail("nextInt() should throw IllegalStateException when hasNext is false");
+ } catch (IllegalStateException e) {
+ // Expected
+ }
+
+ try {
+ reader.isNextNull();
+ Assert.fail("isNextNull() should throw IllegalStateException when hasNext is false");
+ } catch (IllegalStateException e) {
+ // Expected
+ }
+
+ try {
+ reader.skipNext();
+ Assert.fail("skipNext() should throw IllegalStateException when hasNext is false");
+ } catch (IllegalStateException e) {
+ // Expected
+ }
+ } finally {
+ reader.close();
+ }
+ }
+
+ /**
+ * Test rewind with type-specific iteration methods.
+ */
+ @Test
+ public void testRewindWithTypeSpecificMethods()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, LONG_COL);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // First pass: read all values
+ long[] firstPassValues = new long[totalDocs];
+ int idx = 0;
+ while (reader.hasNext()) {
+ if (!reader.isNextNull()) {
+ firstPassValues[idx++] = reader.nextLong();
+ } else {
+ reader.skipNext();
+ idx++;
+ }
+ }
+
+ Assert.assertFalse(reader.hasNext(), "Should be at end after first pass");
+
+ // Rewind
+ reader.rewind();
+ Assert.assertTrue(reader.hasNext(), "Should have values after rewind");
+
+ // Second pass: verify values match
+ idx = 0;
+ while (reader.hasNext()) {
+ if (!reader.isNextNull()) {
+ long value = reader.nextLong();
+ Assert.assertEquals(value, firstPassValues[idx++], "Value mismatch on second pass at index " + idx);
+ } else {
+ reader.skipNext();
+ idx++;
+ }
+ }
+
+ Assert.assertEquals(idx, totalDocs, "Should read same number of documents on second pass");
+ } finally {
+ reader.close();
+ }
+ }
+
+ /**
+ * Test consistency between all three iteration patterns.
+ */
+ @Test
+ public void testAllIterationPatternsConsistency()
+ throws Exception {
+ PinotSegmentColumnReaderImpl reader = new PinotSegmentColumnReaderImpl(_dictEncodedSegment, INT_COL_1);
+
+ try {
+ int totalDocs = reader.getTotalDocs();
+
+ // Pattern 1: Read using next() and null checks
+ reader.rewind();
+ Integer[] pattern1Values = new Integer[totalDocs];
+ int idx = 0;
+ while (reader.hasNext()) {
+ Object value = reader.next();
+ pattern1Values[idx++] = (Integer) value;
+ }
+
+ // Pattern 2: Read using isNextNull() + nextInt()
+ reader.rewind();
+ Integer[] pattern2Values = new Integer[totalDocs];
+ idx = 0;
+ while (reader.hasNext()) {
+ if (!reader.isNextNull()) {
+ pattern2Values[idx++] = reader.nextInt();
+ } else {
+ reader.skipNext();
+ pattern2Values[idx++] = null;
+ }
+ }
+
+ // Pattern 3: Read using getTotalDocs() + getInt(docId)
+ Integer[] pattern3Values = new Integer[totalDocs];
+ for (int docId = 0; docId < totalDocs; docId++) {
+ if (!reader.isNull(docId)) {
+ pattern3Values[docId] = reader.getInt(docId);
+ } else {
+ pattern3Values[docId] = null;
+ }
+ }
+
+ // Compare all patterns
+ for (int i = 0; i < totalDocs; i++) {
+ Assert.assertEquals(pattern1Values[i], pattern2Values[i],
+ "Pattern 1 and Pattern 2 differ at index " + i);
+ Assert.assertEquals(pattern1Values[i], pattern3Values[i],
+ "Pattern 1 and Pattern 3 differ at index " + i);
+ Assert.assertEquals(pattern2Values[i], pattern3Values[i],
+ "Pattern 2 and Pattern 3 differ at index " + i);
+ }
+ } finally {
+ reader.close();
+ }
+ }
+}
diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentCreator.java
index dce1d5b1d4..009f20b231 100644
--- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentCreator.java
+++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentCreator.java
@@ -28,6 +28,7 @@
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.segment.spi.index.creator.SegmentIndexCreationInfo;
import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.ColumnReader;
import org.apache.pinot.spi.data.readers.GenericRow;
@@ -68,6 +69,9 @@ void indexRow(GenericRow row)
void indexColumn(String columnName, @Nullable int[] sortedDocIds, IndexSegment segment)
throws IOException;
+ void indexColumn(String columnName, ColumnReader columnReader)
+ throws IOException;
+
/**
* Sets the name of the segment.
*
diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/IndexCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/IndexCreator.java
index 3e2389f388..105677e195 100644
--- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/IndexCreator.java
+++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/IndexCreator.java
@@ -62,4 +62,89 @@ void add(@Nonnull Object[] values, @Nullable int[] dictIds)
void seal()
throws IOException;
+
+ /**
+ * Primitive type additions for columnar processing optimization.
+ * These methods avoid boxing overhead when iterating over columnar data.
+ * Default implementation boxes the value for backward compatibility.
+ */
+
+ default void addInt(int value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ default void addLong(long value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ default void addFloat(float value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ default void addDouble(double value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ default void addString(String value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ default void addBytes(byte[] value, int dictId)
+ throws IOException {
+ add(value, dictId);
+ }
+
+ // The default implementations box the values for backward compatibility.
+ // This is extremely inefficient because the implementations of add(Object[], int[]) method will end up
+ // unboxing them again to write to the index.
+ default void addIntMV(int[] values, @Nullable int[] dictIds)
+ throws IOException {
+ Integer[] boxedValues = new Integer[values.length];
+ for (int i = 0; i < values.length; i++) {
+ boxedValues[i] = values[i];
+ }
+ add(boxedValues, dictIds);
+ }
+
+ default void addLongMV(long[] values, @Nullable int[] dictIds)
+ throws IOException {
+ Long[] boxedValues = new Long[values.length];
+ for (int i = 0; i < values.length; i++) {
+ boxedValues[i] = values[i];
+ }
+ add(boxedValues, dictIds);
+ }
+
+ default void addFloatMV(float[] values, @Nullable int[] dictIds)
+ throws IOException {
+ Float[] boxedValues = new Float[values.length];
+ for (int i = 0; i < values.length; i++) {
+ boxedValues[i] = values[i];
+ }
+ add(boxedValues, dictIds);
+ }
+
+ default void addDoubleMV(double[] values, @Nullable int[] dictIds)
+ throws IOException {
+ Double[] boxedValues = new Double[values.length];
+ for (int i = 0; i < values.length; i++) {
+ boxedValues[i] = values[i];
+ }
+ add(boxedValues, dictIds);
+ }
+
+ default void addStringMV(String[] values, @Nullable int[] dictIds)
+ throws IOException {
+ add(values, dictIds);
+ }
+
+ default void addBytesMV(byte[][] values, @Nullable int[] dictIds)
+ throws IOException {
+ add(values, dictIds);
+ }
}
diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/CombinedInvertedIndexCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/CombinedInvertedIndexCreator.java
index d75011b14c..9dc8e36654 100644
--- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/CombinedInvertedIndexCreator.java
+++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/CombinedInvertedIndexCreator.java
@@ -103,4 +103,117 @@ default void add(Object[] values, @Nullable int[] dictIds) {
}
}
}
+
+ /**
+ * Primitive type additions for columnar processing optimization.
+ * These methods avoid boxing overhead when iterating over columnar data.
+ */
+
+ @Override
+ default void addInt(int value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ add(value);
+ }
+ }
+
+ @Override
+ default void addLong(long value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ add(value);
+ }
+ }
+
+ @Override
+ default void addFloat(float value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ add(value);
+ }
+ }
+
+ @Override
+ default void addDouble(double value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ add(value);
+ }
+ }
+
+ @Override
+ default void addString(String value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ throw new RuntimeException("String not supported for range index");
+ }
+ }
+
+ @Override
+ default void addBytes(byte[] value, int dictId) {
+ if (dictId >= 0) {
+ add(dictId);
+ } else {
+ throw new RuntimeException("Bytes not supported for range index");
+ }
+ }
+
+ @Override
+ default void addIntMV(int[] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ add(values, values.length);
+ }
+ }
+
+ @Override
+ default void addLongMV(long[] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ add(values, values.length);
+ }
+ }
+
+ @Override
+ default void addFloatMV(float[] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ add(values, values.length);
+ }
+ }
+
+ @Override
+ default void addDoubleMV(double[] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ add(values, values.length);
+ }
+ }
+
+ @Override
+ default void addStringMV(String[] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ throw new RuntimeException("String MV not supported for range index");
+ }
+ }
+
+ @Override
+ default void addBytesMV(byte[][] values, @Nullable int[] dictIds) {
+ if (dictIds != null) {
+ add(dictIds, dictIds.length);
+ } else {
+ throw new RuntimeException("Bytes MV not supported for range index");
+ }
+ }
}
diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/DictionaryBasedInvertedIndexCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/DictionaryBasedInvertedIndexCreator.java
index 96a718ceda..88d10679ff 100644
--- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/DictionaryBasedInvertedIndexCreator.java
+++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/DictionaryBasedInvertedIndexCreator.java
@@ -70,6 +70,71 @@ default void add(@Nonnull Object[] values, @Nullable int[] dictIds) {
add(dictIds, dictIds.length);
}
+ /**
+ * Primitive type additions for columnar processing optimization.
+ * These methods avoid boxing overhead when iterating over columnar data.
+ */
+
+ @Override
+ default void addInt(int value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addLong(long value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addFloat(float value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addDouble(double value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addString(String value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addBytes(byte[] value, int dictId) {
+ add(dictId);
+ }
+
+ @Override
+ default void addIntMV(int[] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
+ @Override
+ default void addLongMV(long[] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
+ @Override
+ default void addFloatMV(float[] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
+ @Override
+ default void addDoubleMV(double[] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
+ @Override
+ default void addStringMV(String[] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
+ @Override
+ default void addBytesMV(byte[][] values, int[] dictIds) {
+ add(dictIds, dictIds.length);
+ }
+
/**
* For single-value column, adds the dictionary id for the next document.
*/
diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java
index dea80c7462..6e92a29c7f 100644
--- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java
+++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java
@@ -151,6 +151,132 @@ default void add(@Nonnull Object[] cellValues, @Nullable int[] dictIds)
}
}
+ /**
+ * Primitive type additions for columnar processing optimization.
+ * These methods avoid boxing overhead when iterating over columnar data.
+ * Default implementations delegate to putDictId or put* methods based on dictId.
+ */
+
+ @Override
+ default void addInt(int value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putInt(value);
+ }
+ }
+
+ @Override
+ default void addLong(long value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putLong(value);
+ }
+ }
+
+ @Override
+ default void addFloat(float value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putFloat(value);
+ }
+ }
+
+ @Override
+ default void addDouble(double value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putDouble(value);
+ }
+ }
+
+ @Override
+ default void addString(String value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putString(value);
+ }
+ }
+
+ @Override
+ default void addBytes(byte[] value, int dictId)
+ throws IOException {
+ if (dictId >= 0) {
+ putDictId(dictId);
+ } else {
+ putBytes(value);
+ }
+ }
+
+ @Override
+ default void addIntMV(int[] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putIntMV(values);
+ }
+ }
+
+ @Override
+ default void addLongMV(long[] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putLongMV(values);
+ }
+ }
+
+ @Override
+ default void addFloatMV(float[] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putFloatMV(values);
+ }
+ }
+
+ @Override
+ default void addDoubleMV(double[] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putDoubleMV(values);
+ }
+ }
+
+ @Override
+ default void addStringMV(String[] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putStringMV(values);
+ }
+ }
+
+ @Override
+ default void addBytesMV(byte[][] values, @Nullable int[] dictIds)
+ throws IOException {
+ if (dictIds != null) {
+ putDictIdMV(dictIds);
+ } else {
+ putBytesMV(values);
+ }
+ }
+
/**
* Returns {@code true} if the forward index is dictionary-encoded, {@code false} if it is raw.
*/
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReader.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReader.java
new file mode 100644
index 0000000000..3c05c26eba
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReader.java
@@ -0,0 +1,354 @@
+/**
+ * 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.spi.data.readers;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.Serializable;
+import javax.annotation.Nullable;
+
+
+/**
+ * The ColumnReader interface is used to read column data from various data sources
+ * for columnar segment building. Unlike RecordReader which reads row-by-row, ColumnReader provides
+ * column-wise access to data, enabling efficient columnar segment creation.
+ *
+ * This interface provides 3 patterns optimised for different use cases
+ * (Some implementations may not support all patterns):
+ *
+ * - Sequential iteration over all values in a column using hasNext(), next() and rewind()
+ * - Sequential iteration with type-specific methods (isInt(), isLong(), nextInt(), nextLong(), etc.)
+ * and null handling (isNextNull(), skipNext()) and supports hasNext() and rewind()
+ * - Random access by document ID using getInt(docId), getLong(docId), etc. and isNull(docId) for null checks
+ *
+ *
+ * Implementations should handle data type conversions and efficient column-wise data access patterns.
+ *
+ *
Usage Patterns
+ *
+ * There are three primary patterns for reading data from a ColumnReader:
+ *
+ *
Pattern 1: Sequential Iteration with Generic next() and Null Checks
+ * This pattern uses the generic {@link #next()} method which returns Object and may return null.
+ * Suitable when you need to handle arbitrary data types or when null handling is done on the return value.
+ *
+ *
{@code
+ * // Read all values in the column
+ * while (columnReader.hasNext()) {
+ * Object value;
+ * try {
+ * value = columnReader.next();
+ * } catch (Exception e) {
+ * // Handle exception / log
+ * continue;
+ * }
+ * if (value != null) {
+ * // Process non-null value
+ * processValue(value);
+ * } else {
+ * // Handle null value
+ * handleNullValue();
+ * }
+ * }
+ *
+ * // Rewind to read the column again
+ * columnReader.rewind();
+ *
+ * // Second pass through the data
+ * while (columnReader.hasNext()) {
+ * Object value = columnReader.next();
+ * if (value != null) {
+ * processValueAgain(value);
+ * }
+ * }
+ * }
+ *
+ * Pattern 2: Sequential Iteration with Type-Specific Methods and Explicit Null Checks
+ * This pattern uses {@link #isNextNull()} to check for nulls before calling type-specific methods
+ * like {@link #nextInt()}, {@link #nextLong()}, etc. Use {@link #skipNext()} to advance past null values.
+ * This is the preferred pattern when you know the column data type and want to avoid boxing overhead.
+ * Before using this pattern, check the column data type using methods like {@link #isInt()}, {@link #isLong()}, etc.
+ * If the data type does not match, fall back to Pattern 1 with {@link #next()}.
+ *
{@code
+ * // Read all int values in the column, handling nulls
+ * if (columnReader.isInt()) {
+ * while (columnReader.hasNext()) {
+ * if (columnReader.isNextNull()) {
+ * // Skip the null value
+ * columnReader.skipNext();
+ * handleNullValue();
+ * } else {
+ * // Read the primitive int value (no boxing)
+ * int value = columnReader.nextInt();
+ * processIntValue(value);
+ * }
+ * }
+ * } else {
+ * // Fallback to Pattern 1 if not INT type
+ * }
+ *
+ * }
+ *
+ * Pattern 3: Random Access by Document ID
+ * This pattern uses {@link #getTotalDocs()} to get the total number of documents, then uses
+ * document ID-based accessors like {@link #getInt(int)}, {@link #getLong(int)}, etc. to read
+ * specific values. Use {@link #isNull(int)} to check if a value is null before reading.
+ * This pattern is useful when you need random access or want to process documents in a specific order.
+ *
+ *
{@code
+ *
+ * // Random access example - read specific document IDs
+ * int[] docIdsToRead = {5, 10, 15, 20};
+ * for (int docId : docIdsToRead) {
+ * if (!columnReader.isNull(docId)) {
+ * int value = columnReader.getInt(docId);
+ * processSpecificDoc(docId, value);
+ * }
+ * }
+ *
+ * // Read in reverse order
+ * // Get the total number of documents
+ * int totalDocs = columnReader.getTotalDocs();
+ * for (int docId = totalDocs - 1; docId >= 0; docId--) {
+ * if (!columnReader.isNull(docId)) {
+ * int value = columnReader.getInt(docId);
+ * processReverseOrder(docId, value);
+ * }
+ * }
+ * }
+ *
+ * Choosing the Right Pattern
+ *
+ * - Pattern 1: Use when dealing with generic Object types or when you don't know
+ * the column type at compile time. Less efficient due to boxing.
+ * - Pattern 2: Use when you know the column type and want efficient sequential iteration
+ * with primitive types. Preferred for most columnar segment building scenarios.
+ * - Pattern 3: Use when you need random access, want to process documents in a specific
+ * order, or need to access the same document multiple times.
+ *
+ *
+ */
+public interface ColumnReader extends Closeable, Serializable {
+
+ /**
+ * Return true if more values remain to be read in this column.
+ * This method should not throw exception. Caller is not responsible for handling exceptions from this method.
+ */
+ boolean hasNext();
+
+ /**
+ * Get the next value in the column.
+ *
This method should be called only if {@link #hasNext()} returns true. Caller is responsible for
+ * handling exceptions from this method and skip the value if user wants to continue reading the remaining values.
+ *
+ * @return Next column value, or null if the value is null
+ * @throws IOException If an I/O error occurs while reading
+ */
+ @Nullable
+ Object next()
+ throws IOException;
+
+ /**
+ * Check if the next value to be read is null.
+ */
+ boolean isNextNull() throws IOException;
+
+ /**
+ * Move the reader to skip the next value in the column and advance to the following value.
+ * This method is typically used because type specific methods like {@link #nextInt()}, {@link #nextLong()}, etc.
+ * cannot return null values
+ * Thus, if {@link #isNextNull()} returns true, clients should call this method to skip the null value.
+ *
+ *
Example
+ *
{@code
+ * if (columnReader.isNextNull()) {
+ * columnReader.skipNext(); // Skip null and move to next value
+ * } else {
+ * int value = columnReader.nextInt();
+ * }
+ * }
+ *
+ * @throws IOException If an I/O error occurs while skipping
+ */
+ void skipNext() throws IOException;
+
+ /**
+ * Check if the column data is single-value or multi-value.
+ */
+ boolean isSingleValue();
+
+ /**
+ * Check if the column data type from the actual reader can be returned as the expected type directly.
+ * For multi-value columns, this indicates if the multi-value type specific methods can be called directly.
+ * If true, the type specific methods like nextInt() can be called directly.
+ * Otherwise, clients should use next() and cast the result.
+ */
+ boolean isInt();
+
+ boolean isLong();
+
+ boolean isFloat();
+
+ boolean isDouble();
+
+ boolean isString();
+
+ boolean isBytes();
+
+ /**
+ * Get the next int / long / float / double / string / byte[] value for single-value columns.
+ * Should be called only if isNextNull() returns false.
+ * @throws IOException If an I/O error occurs while reading
+ */
+ int nextInt() throws IOException;
+
+ long nextLong() throws IOException;
+
+ float nextFloat() throws IOException;
+
+ double nextDouble() throws IOException;
+
+ String nextString() throws IOException;
+
+ byte[] nextBytes() throws IOException;
+
+ /**
+ * Get the next int[] / long[] / float[] / double[] / string[] / bytes[][] values for multi-value columns.
+ * Should be called only if isNextNull() returns false.
+ *
+ * For primitive types (int, long, float, double), returns a {@link MultiValueResult} that includes
+ * element-level null validity tracking. Use {@link MultiValueResult#hasNulls()} and
+ * {@link MultiValueResult#isNull(int)} to check for null elements within the array.
+ *
+ * @throws IOException If an I/O error occurs while reading
+ */
+ MultiValueResult nextIntMV()
+ throws IOException;
+
+ MultiValueResult nextLongMV()
+ throws IOException;
+
+ MultiValueResult nextFloatMV()
+ throws IOException;
+
+ MultiValueResult nextDoubleMV()
+ throws IOException;
+
+ String[] nextStringMV()
+ throws IOException;
+
+ byte[][] nextBytesMV()
+ throws IOException;
+
+ /**
+ * Rewind the reader to start reading from the first value again.
+ *
+ * @throws IOException If an I/O error occurs while rewinding
+ */
+ void rewind() throws IOException;
+
+ /**
+ * Get the name of the column.
+ *
+ * @return Column name
+ */
+ String getColumnName();
+
+ /**
+ * Get the total number of documents in this column.
+ *
+ * @return Total number of documents
+ */
+ int getTotalDocs();
+
+ /**
+ * Check if the value at the given document ID is null.
+ * Document ID is 0-based. Valid values are 0 to {@link #getTotalDocs()} - 1.
+ *
+ * @param docId Document ID (0-based)
+ * @return true if the value is null, false otherwise
+ * @throws IndexOutOfBoundsException If docId is out of range
+ */
+ boolean isNull(int docId) throws IOException;
+
+ // Single-value accessors
+
+ /**
+ * Get int / long / float / double / string / byte[] value at the given document ID for single-value columns.
+ * Should be called only if isNull(docId) returns false.
+ * Document ID is 0-based. Valid values are 0 to {@link #getTotalDocs()} - 1.
+ *
+ * @param docId Document ID (0-based)
+ * @throws IndexOutOfBoundsException If docId is out of range
+ * @throws IOException If an I/O error occurs while reading
+ */
+ int getInt(int docId) throws IOException;
+
+ long getLong(int docId) throws IOException;
+
+ float getFloat(int docId) throws IOException;
+
+ double getDouble(int docId) throws IOException;
+
+ String getString(int docId) throws IOException;
+
+ byte[] getBytes(int docId) throws IOException;
+
+ /**
+ * Get the value at the given document ID as a Java Object.
+ * Can be used for both single-value and multi-value columns.
+ * This should be used if
+ * 1. Certain API's don't yet support primitive type specific methods (eg: TimeHandler, Partitioner, etc.) and
+ * thus will be boxed anyway.
+ * 2. The required data type does not match the actual type and the client will handle the conversion
+ * Document ID is 0-based. Valid values are 0 to {@link #getTotalDocs()} - 1.
+ *
+ * @param docId Document ID (0-based)
+ * @throws IndexOutOfBoundsException If docId is out of range
+ * @throws IOException If an I/O error occurs while reading
+ */
+ Object getValue(int docId) throws IOException;
+
+ // Multi-value accessors
+
+ /**
+ * Get int[] / long[] / float[] / double[] / string[] / bytes[][] values at the given doc ID for multi-value columns.
+ * Should be called only if isNull(docId) returns false.
+ *
Document ID is 0-based. Valid values are 0 to {@link #getTotalDocs()} - 1.
+ *
+ *
For primitive types (int, long, float, double), returns a {@link MultiValueResult} that includes
+ * element-level null validity tracking. Use {@link MultiValueResult#hasNulls()} and
+ * {@link MultiValueResult#isNull(int)} to check for null elements within the array.
+ *
+ * @param docId Document ID (0-based)
+ * @throws IndexOutOfBoundsException If docId is out of range
+ * @throws IOException If an I/O error occurs while reading
+ */
+ MultiValueResult getIntMV(int docId) throws IOException;
+
+ MultiValueResult getLongMV(int docId) throws IOException;
+
+ MultiValueResult getFloatMV(int docId) throws IOException;
+
+ MultiValueResult getDoubleMV(int docId) throws IOException;
+
+ String[] getStringMV(int docId) throws IOException;
+
+ byte[][] getBytesMV(int docId) throws IOException;
+}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReaderFactory.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReaderFactory.java
new file mode 100644
index 0000000000..ce9a9c3cc9
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReaderFactory.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.spi.data.readers;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.data.Schema;
+
+
+/**
+ * Factory interface for creating ColumnReader instances for different data sources.
+ *
+ * This factory provides a generic way to create column readers for various data sources
+ * such as Pinot segments, Parquet files, etc. The factory handles:
+ *
+ * - Creating column readers for existing columns
+ * - Creating default value readers for new columns
+ * - Data type conversions between source and target schemas
+ * - Resource management
+ *
+ */
+public interface ColumnReaderFactory extends Closeable, Serializable {
+
+ /**
+ * Configuration key to enable random access mode for column readers.
+ * When set to "true", implementations may cache data to support non-sequential access patterns.
+ */
+ String CONFIG_ENABLE_RANDOM_ACCESS = "enableRandomAccess";
+
+ /**
+ * Initialize the factory with the data source and target schema.
+ *
+ * @param targetSchema Target schema for the output segment
+ * @throws IOException If initialization fails
+ */
+ void init(Schema targetSchema)
+ throws IOException;
+
+ /**
+ * Initialize the factory with the data source, target schema, and specific target columns.
+ * @param targetSchema Target schema for the output segment
+ * @param colsToRead Set of target columns to read
+ * @throws IOException If initialization fails
+ */
+ void init(Schema targetSchema, Set colsToRead)
+ throws IOException;
+
+ /**
+ * Initialize the factory with the data source, target schema, specific columns, and configuration.
+ * @param targetSchema Target schema for the output segment
+ * @param colsToRead Set of target columns to read
+ * @param configs Configuration map for implementation-specific settings
+ * @throws IOException If initialization fails
+ */
+ default void init(Schema targetSchema, Set colsToRead, Map configs)
+ throws IOException {
+ init(targetSchema, colsToRead);
+ }
+
+ /**
+ * Get the set of column names available in the source data.
+ *
+ * @return Set of available column names in the source
+ */
+ Set getAvailableColumns();
+
+ /**
+ * Get a column reader for the specified column.
+ * Implementations may cache and reuse readers for efficiency.
+ *
+ * @param columnName Name of the column to read
+ * Can return null if column doesn't exist in the source
+ * @return ColumnReader instance for the specified column (may be cached)
+ */
+ @Nullable ColumnReader getColumnReader(String columnName);
+
+ /**
+ * Get all column readers for the target schema.
+ * Implementations may cache and reuse instances that were created during init().
+ *
+ * @return Map of column name to ColumnReader (cached instances)
+ * @throws IOException If any column reader cannot be obtained
+ */
+ Map getAllColumnReaders()
+ throws IOException;
+}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnarDataSource.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnarDataSource.java
new file mode 100644
index 0000000000..a0d1da64cc
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnarDataSource.java
@@ -0,0 +1,66 @@
+/**
+ * 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.spi.data.readers;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+/**
+ *
+ * This interface is designed to support the creation of multiple {@link ColumnReaderFactory} instances
+ * from the same underlying data source. This capability is required because different processing contexts
+ * may require distinct column reader objects for the same input file and column. For example:
+ *
+ * - Different processing stages may need independent iterators over the same column
+ * - Parallel processing threads may each require their own column readers
+ *
+ *
+ */
+public interface ColumnarDataSource extends Closeable {
+
+ /**
+ * Returns the total number of documents (rows) in this data source.
+ *
+ * @return the total document count
+ */
+ int getTotalDocs();
+
+ /**
+ * Creates a new {@link ColumnReaderFactory} instance for this data source.
+ *
+ * Multiple factories can be created from the same data source, allowing different consumers
+ * to independently read and process the columnar data. Each factory can then create its own
+ * set of column readers without interfering with readers created by other factories.
+ *
+ * @return a new column reader factory instance
+ * @throws IOException if an I/O error occurs while creating the factory
+ */
+ ColumnReaderFactory createColumnReaderFactory()
+ throws IOException;
+
+ /**
+ * Returns a string description of the underlying data source.
+ *
+ * This typically includes information such as the file name, path, or other identifiers
+ * that help identify the source of the data.
+ *
+ * @return a description of the underlying source (e.g., file name)
+ */
+ String toString();
+}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/MultiValueResult.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/MultiValueResult.java
new file mode 100644
index 0000000000..3517f7ba02
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/MultiValueResult.java
@@ -0,0 +1,123 @@
+/**
+ * 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.spi.data.readers;
+
+import java.util.BitSet;
+import javax.annotation.Nullable;
+
+
+/**
+ * Result wrapper for multi-value column reads that tracks element-level nulls.
+ *
+ *
This class addresses a limitation where bulk reads from columnar formats (like Arrow)
+ * don't check the validity bitmap for null elements. By returning both the values array
+ * and a nulls BitSet, callers can properly handle null elements within multi-value columns.
+ *
+ *
BitSet Semantics:
+ *
+ * - Set bit (1) = null element
+ * - Unset bit (0) = valid/non-null element
+ * - Null nulls BitSet = no nulls in the range (fast path)
+ *
+ *
+ * Usage Example:
+ *
{@code
+ * MultiValueResult result = columnReader.nextIntMV();
+ * int[] values = result.getValues();
+ *
+ * if (result.hasNulls()) {
+ * for (int i = 0; i < values.length; i++) {
+ * if (result.isNull(i)) {
+ * // Handle null element - values[i] contains default value (0 for int)
+ * } else {
+ * // Use values[i]
+ * }
+ * }
+ * } else {
+ * // Fast path: no nulls, use values array directly
+ * }
+ * }
+ *
+ * @param The array type (int[], long[], float[], double[])
+ */
+public class MultiValueResult {
+ private final T _values;
+ @Nullable
+ private final BitSet _nulls;
+
+ private MultiValueResult(T values, @Nullable BitSet nulls) {
+ _values = values;
+ _nulls = nulls;
+ }
+
+ /**
+ * Create a MultiValueResult with optional nulls information.
+ *
+ * @param values The values array (int[], long[], float[], double[])
+ * @param nulls BitSet where set bits indicate null elements,
+ * or null if no nulls exist in the range
+ * @param The array type
+ * @return A new MultiValueResult instance
+ */
+ public static MultiValueResult of(T values, @Nullable BitSet nulls) {
+ return new MultiValueResult<>(values, nulls);
+ }
+
+ /**
+ * Check if any elements in this result are null.
+ *
+ * @return true if there are null elements, false otherwise
+ */
+ public boolean hasNulls() {
+ return _nulls != null;
+ }
+
+ /**
+ * Check if a specific element is null.
+ *
+ * @param index The index of the element to check
+ * @return true if the element is null, false if it's valid
+ */
+ public boolean isNull(int index) {
+ return _nulls != null && _nulls.get(index);
+ }
+
+ /**
+ * Get the values array.
+ *
+ * Note: If {@link #hasNulls()} returns true, some elements in this array
+ * may contain default values (0 for numeric types) for null positions.
+ * Use {@link #isNull(int)} to check individual elements.
+ *
+ * @return The values array
+ */
+ public T getValues() {
+ return _values;
+ }
+
+ /**
+ * Get the nulls BitSet.
+ *
+ * @return The nulls BitSet where set bits indicate null elements, or null if no nulls exist
+ */
+ @Nullable
+ public BitSet getNulls() {
+ return _nulls;
+ }
+}