From 33df05d9add6b214cc31b7289734de20b4c403bb Mon Sep 17 00:00:00 2001 From: rohit Date: Tue, 4 Nov 2025 14:01:29 +0530 Subject: [PATCH 1/8] column major segment build for columnar datasource (#16727) * build segment columnarly from a column major datasource * resolve review comments 1/2 * resolve review comments 2/2 * always close column readers * refactor ColumnReaderFactory * resolved all review comments (cherry picked from commit 13da6e3935d183b235dc377a930093bf16a89952) --- .../ColumnarSegmentCreationDataSource.java | 95 ++++ .../impl/SegmentColumnarIndexCreator.java | 48 ++ .../impl/SegmentIndexCreationDriverImpl.java | 169 +++++- ...ColumnarSegmentPreIndexStatsContainer.java | 178 +++++++ .../SegmentPreIndexStatsCollectorImpl.java | 32 +- .../impl/stats/StatsCollectorUtil.java | 65 +++ .../readers/DefaultValueColumnReader.java | 94 ++++ .../PinotSegmentColumnReaderFactory.java | 181 +++++++ .../readers/PinotSegmentColumnReaderImpl.java | 106 ++++ .../impl/ColumnReaderInterfaceTest.java | 129 +++++ .../impl/ColumnarRowMajorEquivalenceTest.java | 85 +++ .../impl/ColumnarSchemaEvolutionTest.java | 55 ++ .../impl/ColumnarSegmentBuildingTestBase.java | 488 ++++++++++++++++++ .../segment/spi/creator/SegmentCreator.java | 4 + .../pinot/spi/data/readers/ColumnReader.java | 76 +++ .../spi/data/readers/ColumnReaderFactory.java | 79 +++ 16 files changed, 1834 insertions(+), 50 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/ColumnarSegmentCreationDataSource.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnarSegmentPreIndexStatsContainer.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StatsCollectorUtil.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderFactory.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnReaderInterfaceTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarRowMajorEquivalenceTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSchemaEvolutionTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSegmentBuildingTestBase.java create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReader.java create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReaderFactory.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/ColumnarSegmentCreationDataSource.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/ColumnarSegmentCreationDataSource.java new file mode 100644 index 0000000000..0dee1fbdca --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/ColumnarSegmentCreationDataSource.java @@ -0,0 +1,95 @@ +/** + * 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; + +import java.io.Closeable; +import java.io.IOException; +import java.util.Map; +import org.apache.pinot.segment.local.segment.creator.impl.stats.ColumnarSegmentPreIndexStatsContainer; +import org.apache.pinot.segment.spi.creator.SegmentCreationDataSource; +import org.apache.pinot.segment.spi.creator.SegmentPreIndexStatsCollector; +import org.apache.pinot.segment.spi.creator.StatsCollectorConfig; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * SegmentCreationDataSource implementation that uses ColumnReaderFactory for columnar data access. + * + *

This data source enables columnar segment building by providing access to column readers + * instead of row-based record readers. It supports: + *

+ */ +public class ColumnarSegmentCreationDataSource implements SegmentCreationDataSource, Closeable { + private static final Logger LOGGER = LoggerFactory.getLogger(ColumnarSegmentCreationDataSource.class); + + private final Map _columnReaders; + + /** + * Create a ColumnarSegmentCreationDataSource. + * + * @param columnReaders Map of column name to ColumnReader instances + */ + public ColumnarSegmentCreationDataSource(Map columnReaders) { + _columnReaders = columnReaders; + } + + @Override + public SegmentPreIndexStatsCollector gatherStats(StatsCollectorConfig statsCollectorConfig) { + LOGGER.info("Gathering stats using columnar approach for {} columns", _columnReaders.size()); + + // Use columnar stats container that efficiently collects statistics column-wise + ColumnarSegmentPreIndexStatsContainer columnarSegmentPreIndexStatsContainer = + new ColumnarSegmentPreIndexStatsContainer(_columnReaders, statsCollectorConfig); + columnarSegmentPreIndexStatsContainer.init(); + return columnarSegmentPreIndexStatsContainer; + } + + @Override + public RecordReader getRecordReader() { + // For columnar building, we don't use RecordReader + // This method is required by the interface but should not be called in columnar mode + throw new UnsupportedOperationException( + "RecordReader not supported in columnar mode. Use getColumnReaders() instead."); + } + + /** + * Get all column readers. + * + * @return Map of column name to ColumnReader + */ + public Map getColumnReaders() { + return _columnReaders; + } + + @Override + public void close() + throws IOException { + for (ColumnReader columnReader : _columnReaders.values()) { + columnReader.close(); + } + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java index a55ef41dfd..eb59cbc7cc 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java @@ -68,6 +68,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.FieldSpec.FieldType; 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.apache.pinot.spi.env.CommonsConfigurationUtils; import org.apache.pinot.spi.utils.TimeUtils; @@ -421,6 +422,53 @@ public void indexColumn(String columnName, @Nullable int[] sortedDocIds, IndexSe } } + /** + * Index a column using a ColumnReader (column-major approach). + * This method processes the column values using the iterator pattern from ColumnReader. + * + * @param columnName Name of the column to index + * @param columnReader ColumnReader for the column data + * @throws IOException if indexing fails + */ + @Override + public void indexColumn(String columnName, ColumnReader columnReader) + throws IOException { + Map, IndexCreator> creatorsByIndex = _creatorsByColAndIndex.get(columnName); + NullValueVectorCreator nullVec = _nullValueVectorCreatorMap.get(columnName); + FieldSpec fieldSpec = _schema.getFieldSpecFor(columnName); + SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); + Object defaultNullValue = fieldSpec.getDefaultNullValue(); + Object reuseColumnValueToIndex; + + // Reset column reader to start from beginning + columnReader.rewind(); + + int docId = 0; + while (columnReader.hasNext()) { + reuseColumnValueToIndex = columnReader.next(); + + // Handle null values + if (reuseColumnValueToIndex == null) { + if (nullVec != null) { + nullVec.setNull(docId); + } + reuseColumnValueToIndex = defaultNullValue; + } + + try { + if (fieldSpec.isSingleValueField()) { + indexSingleValueRow(dictionaryCreator, reuseColumnValueToIndex, creatorsByIndex); + } else { + indexMultiValueRow(dictionaryCreator, (Object[]) reuseColumnValueToIndex, creatorsByIndex); + } + } catch (JsonParseException jpe) { + throw new ColumnJsonParserException(columnName, jpe); + } + + docId++; + } + } + private void indexColumnValue(PinotSegmentColumnReader colReader, Map, IndexCreator> creatorsByIndex, String columnName, FieldSpec fieldSpec, SegmentDictionaryCreator dictionaryCreator, int sourceDocId, int onDiskDocPos, diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java index 612b696e27..a7c8857a7c 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java @@ -36,6 +36,7 @@ import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.realtime.converter.stats.RealtimeSegmentSegmentCreationDataSource; +import org.apache.pinot.segment.local.segment.creator.ColumnarSegmentCreationDataSource; import org.apache.pinot.segment.local.segment.creator.RecordReaderSegmentCreationDataSource; import org.apache.pinot.segment.local.segment.creator.TransformPipeline; import org.apache.pinot.segment.local.segment.index.converter.SegmentFormatConverterFactory; @@ -75,6 +76,8 @@ import org.apache.pinot.spi.data.IngestionSchemaValidator; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.SchemaValidatorFactory; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.ColumnReaderFactory; import org.apache.pinot.spi.data.readers.FileFormat; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.data.readers.RecordReader; @@ -84,8 +87,6 @@ import org.apache.pinot.spi.utils.ReadMode; import org.slf4j.Logger; import org.slf4j.LoggerFactory; - - /** * Implementation of an index segment creator. */ @@ -94,7 +95,7 @@ public class SegmentIndexCreationDriverImpl implements SegmentIndexCreationDrive private static final Logger LOGGER = LoggerFactory.getLogger(SegmentIndexCreationDriverImpl.class); private SegmentGeneratorConfig _config; - private RecordReader _recordReader; + @Nullable private RecordReader _recordReader; private SegmentPreIndexStatsContainer _segmentStats; // NOTE: Use TreeMap so that the columns are ordered alphabetically private TreeMap _indexCreationInfoMap; @@ -102,7 +103,7 @@ public class SegmentIndexCreationDriverImpl implements SegmentIndexCreationDrive private SegmentIndexCreationInfo _segmentIndexCreationInfo; private SegmentCreationDataSource _dataSource; private Schema _dataSchema; - private TransformPipeline _transformPipeline; + @Nullable private TransformPipeline _transformPipeline; private IngestionSchemaValidator _ingestionSchemaValidator; private int _totalDocs = 0; private File _tempIndexDir; @@ -162,31 +163,71 @@ public void init(SegmentGeneratorConfig config, RecordReader recordReader) new TransformPipeline(config.getTableConfig(), config.getSchema())); } + /** + * Initialize the driver for columnar segment building using a ColumnReaderFactory. + * This method sets up the driver to use column-wise input data access instead of row-wise. + * + * @param config Segment generator configuration + * @param columnReaderFactory Factory for creating column readers + * @throws Exception if initialization fails + */ + public void init(SegmentGeneratorConfig config, ColumnReaderFactory columnReaderFactory) + throws Exception { + // Initialize the column reader factory with target schema + columnReaderFactory.init(config.getSchema()); + + // Get all column readers for the target schema + Map columnReaders = columnReaderFactory.getAllColumnReaders(); + + // Create columnar data source + ColumnarSegmentCreationDataSource columnarDataSource = new ColumnarSegmentCreationDataSource(columnReaders); + + // Use the existing init method with columnar data source and no transform pipeline + init(config, columnarDataSource, null); + + LOGGER.info("Initialized SegmentIndexCreationDriverImpl for columnar data source building with {} columns", + columnReaders.size()); + } + public void init(SegmentGeneratorConfig config, SegmentCreationDataSource dataSource, TransformPipeline transformPipeline) throws Exception { _config = config; - _recordReader = dataSource.getRecordReader(); _dataSchema = config.getSchema(); _continueOnError = config.isContinueOnError(); + String readerClassName = null; + + // Handle columnar data sources differently + if (dataSource instanceof ColumnarSegmentCreationDataSource) { + // For columnar data sources, we don't have a record reader + _recordReader = null; + _transformPipeline = null; // No transform pipeline for columnar mode + _dataSource = dataSource; + } else { + // For record reader-based data sources + _recordReader = dataSource.getRecordReader(); - if (config.isFailOnEmptySegment()) { - Preconditions.checkState(_recordReader.hasNext(), "No record in data source"); - } - _transformPipeline = transformPipeline; - // Use the same transform pipeline if the data source is backed by a record reader - if (dataSource instanceof RecordReaderSegmentCreationDataSource) { - ((RecordReaderSegmentCreationDataSource) dataSource).setTransformPipeline(transformPipeline); - } + if (config.isFailOnEmptySegment()) { + Preconditions.checkState(_recordReader.hasNext(), "No record in data source"); + } + _transformPipeline = transformPipeline; + // Use the same transform pipeline if the data source is backed by a record reader + if (dataSource instanceof RecordReaderSegmentCreationDataSource) { + ((RecordReaderSegmentCreationDataSource) dataSource).setTransformPipeline(transformPipeline); + } - // Optimization for realtime segment conversion - if (dataSource instanceof RealtimeSegmentSegmentCreationDataSource) { - _config.setRealtimeConversion(true); - _config.setConsumerDir(((RealtimeSegmentSegmentCreationDataSource) dataSource).getConsumerDir()); - } + // Optimization for realtime segment conversion + if (dataSource instanceof RealtimeSegmentSegmentCreationDataSource) { + _config.setRealtimeConversion(true); + _config.setConsumerDir(((RealtimeSegmentSegmentCreationDataSource) dataSource).getConsumerDir()); + } - // For stats collection - _dataSource = dataSource; + // For stats collection + _dataSource = dataSource; + + // Initialize common components + readerClassName = _recordReader.getClass().getName(); + } // Initialize index creation _segmentIndexCreationInfo = new SegmentIndexCreationInfo(); @@ -202,7 +243,7 @@ public void init(SegmentGeneratorConfig config, SegmentCreationDataSource dataSo } _ingestionSchemaValidator = - SchemaValidatorFactory.getSchemaValidator(_dataSchema, _recordReader.getClass().getName(), + SchemaValidatorFactory.getSchemaValidator(_dataSchema, readerClassName, config.getInputFilePath()); // Create a temporary directory used in segment creation @@ -230,6 +271,13 @@ private int[] getImmutableToMutableIdMap(@Nullable int[] sortedDocIds) { @Override public void build() throws Exception { + // Check if we're using a columnar data source and switch to columnar mode + if (_dataSource instanceof ColumnarSegmentCreationDataSource) { + LOGGER.info("Detected columnar data source, using columnar building approach"); + buildColumnar(); + return; + } + // Count the number of documents and gather per-column statistics LOGGER.debug("Start building StatsCollector!"); collectStatsAndIndexCreationInfo(); @@ -586,4 +634,83 @@ public IngestionSchemaValidator getIngestionSchemaValidator() { public SegmentPreIndexStatsContainer getSegmentStats() { return _segmentStats; } + + /** + * Build segment using columnar approach. + * This method builds the segment by processing data column-wise instead of row-wise. + * Following is not supported: + *
  • recort transformation + *
  • 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/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/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..22a001ae3b --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java @@ -0,0 +1,94 @@ +/** + * 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.spi.data.FieldSpec; +import org.apache.pinot.spi.data.readers.ColumnReader; + + +/** + * 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 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; + + // 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}; + } + } + + @Override + public boolean hasNext() { + return _currentIndex < _numDocs; + } + + @Override + @Nullable + public Object next() + throws IOException { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + _currentIndex++; + return _defaultValue; + } + + @Override + public void rewind() + throws IOException { + _currentIndex = 0; + } + + @Override + public String getColumnName() { + return _columnName; + } + + @Override + public void close() + throws IOException { + // No resources to close + } +} 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..a69af20d8a --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderFactory.java @@ -0,0 +1,181 @@ +/** + * 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; + @Nullable private Map _columnReaders; + + /** + * Create a PinotSegmentColumnReaderFactory. + * + * @param indexSegment Source segment to read from + */ + public PinotSegmentColumnReaderFactory(IndexSegment indexSegment) { + _indexSegment = indexSegment; + _columnReaders = new HashMap<>(); + } + + @Override + public void init(Schema targetSchema) + throws IOException { + _targetSchema = targetSchema; + _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) + throws IOException { + if (_targetSchema == null) { + throw new IllegalStateException("Factory not initialized. Call init() first."); + } + + ColumnReader reader = _columnReaders.get(columnName); + if (reader == null) { + throw new IOException("Column reader not found for column: " + columnName); + } + return reader; + } + + /** + * Internal method to create a column reader for the specified column. + * This method is called during initialization to create all readers. + */ + private ColumnReader createColumnReader(String columnName, FieldSpec targetFieldSpec) { + if (targetFieldSpec.isVirtualColumn()) { + throw new IllegalStateException("Target field spec is a virtual column."); + } + + ColumnReader columnReader; + + 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 { + // 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() + throws IOException { + 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 (FieldSpec fieldSpec : _targetSchema.getAllFieldSpecs()) { + if (fieldSpec.isVirtualColumn()) { + continue; + } + + String columnName = fieldSpec.getName(); + ColumnReader reader = createColumnReader(columnName, fieldSpec); + 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..9a09394466 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImpl.java @@ -0,0 +1,106 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.segment.readers; + +import java.io.IOException; +import javax.annotation.Nullable; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.spi.data.readers.ColumnReader; + + +/** + * 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 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; + } + + @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 void rewind() + throws IOException { + _currentIndex = 0; + _reuseValue = null; + } + + @Override + public String getColumnName() { + return _columnName; + } + + @Override + public void close() + throws IOException { + _segmentColumnReader.close(); + } +} 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..f65afe106e --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarSegmentBuildingTestBase.java @@ -0,0 +1,488 @@ +/** + * 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_STRING_COL = "mvStringCol"; + + // 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_STRING_COL, FieldSpec.DataType.STRING) + .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_STRING_COL, FieldSpec.DataType.STRING) + .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 { + File outputDir = new File(_tempDir, "rowMajorSegment"); + FileUtils.deleteQuietly(outputDir); + outputDir.mkdirs(); + + SegmentGeneratorConfig config = new SegmentGeneratorConfig(_tableConfig, _originalSchema); + config.setOutDir(outputDir.getAbsolutePath()); + config.setSegmentName(SEGMENT_NAME + "_rowMajor"); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + TestRecordReader recordReader = new TestRecordReader(_testData); + driver.init(config, recordReader); + driver.build(); + + return new File(outputDir, SEGMENT_NAME + "_rowMajor"); + } + + 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_STRING_COL, random.nextDouble() < nullProbability ? null + : new Object[]{"mv1_" + i, "mv2_" + i, "mv3_" + (i % 3)}); + + 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-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-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..ce3f898a6e --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReader.java @@ -0,0 +1,76 @@ +/** + * 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 follows an iterator pattern similar to RecordReader: + *

      + *
    • Sequential iteration over all values in a column using hasNext() and next()
      • + *
    • Rewind capability to restart iteration
      • + *
    • Resource cleanup
      • + *
    + * + *

    Implementations should handle data type conversions, default values for new columns, + * and efficient column-wise data access patterns. + */ +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; + + /** + * 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(); +} 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..1df2be955c --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnReaderFactory.java @@ -0,0 +1,79 @@ +/** + * 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 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 { + + /** + * 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; + + /** + * 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 + * @return ColumnReader instance for the specified column (may be cached) + * @throws IOException If the column reader cannot be obtained + */ + ColumnReader getColumnReader(String columnName) + throws IOException; + + /** + * 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; +} From 13f6b3c9e96839f4580733a3ece7c800923a32d5 Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Tue, 18 Nov 2025 18:03:37 +0530 Subject: [PATCH 2/8] Add APIs for primitive data types for stats collectors, index creators, column readers (#17183) * Add APIs for primitive data types for stats collectors, index creators, column readers * checkstyle fixes * Simplify assertions for index out of bounds check * Add nextInt(), etc methods in ColumnReader * Update documentation and add few more methods like isInt() etc to ColumnReader (cherry picked from commit 560858552e4947b16b9acf2fb572f9104abd7459) --- .../impl/SegmentDictionaryCreator.java | 90 ++ .../AbstractColumnStatisticsCollector.java | 34 + .../DoubleColumnPreIndexStatsCollector.java | 48 +- .../FloatColumnPreIndexStatsCollector.java | 48 +- .../IntColumnPreIndexStatsCollector.java | 48 +- .../LongColumnPreIndexStatsCollector.java | 48 +- .../readers/DefaultValueColumnReader.java | 316 +++++- .../readers/PinotSegmentColumnReader.java | 118 +++ .../PinotSegmentColumnReaderFactory.java | 19 +- .../readers/PinotSegmentColumnReaderImpl.java | 261 +++++ .../impl/ColumnarSegmentBuildingTestBase.java | 98 +- ...AbstractColumnStatisticsCollectorTest.java | 708 +++++++++++++ .../readers/DefaultValueColumnReaderTest.java | 613 ++++++++++++ .../PinotSegmentColumnReaderImplTest.java | 927 ++++++++++++++++++ .../pinot/segment/spi/index/IndexCreator.java | 85 ++ .../creator/CombinedInvertedIndexCreator.java | 113 +++ .../DictionaryBasedInvertedIndexCreator.java | 65 ++ .../index/creator/ForwardIndexCreator.java | 126 +++ .../pinot/spi/data/readers/ColumnReader.java | 235 ++++- .../spi/data/readers/ColumnReaderFactory.java | 9 + 20 files changed, 3899 insertions(+), 110 deletions(-) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollectorTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReaderTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java index 2e4db2a6e6..b27f39ca69 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java @@ -332,6 +332,96 @@ public int indexOfSV(Object value) { } } + /** + * Get dictionary index for a primitive int value without boxing. + */ + public int indexOfSV(int value) { + return _intValueToIndexMap.get(value); + } + + /** + * Get dictionary index for a primitive long value without boxing. + */ + public int indexOfSV(long value) { + return _longValueToIndexMap.get(value); + } + + /** + * Get dictionary index for a primitive float value without boxing. + */ + public int indexOfSV(float value) { + return _floatValueToIndexMap.get(value); + } + + /** + * Get dictionary index for a primitive double value without boxing. + */ + public int indexOfSV(double value) { + return _doubleValueToIndexMap.get(value); + } + + /** + * Get dictionary index for a String value. + */ + public int indexOfSV(String value) { + return _objectValueToIndexMap.getInt(value); + } + + /** + * Get dictionary index for a byte array value. + */ + public int indexOfSV(byte[] value) { + return _objectValueToIndexMap.getInt(new ByteArray(value)); + } + + public int[] indexOfMV(int[] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _intValueToIndexMap.get(values[i]); + } + return indexes; + } + + public int[] indexOfMV(long[] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _longValueToIndexMap.get(values[i]); + } + return indexes; + } + + public int[] indexOfMV(float[] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _floatValueToIndexMap.get(values[i]); + } + return indexes; + } + + public int[] indexOfMV(double[] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _doubleValueToIndexMap.get(values[i]); + } + return indexes; + } + + public int[] indexOfMV(String[] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _objectValueToIndexMap.getInt(values[i]); + } + return indexes; + } + + public int[] indexOfMV(byte[][] values) { + int[] indexes = new int[values.length]; + for (int i = 0; i < values.length; i++) { + indexes[i] = _objectValueToIndexMap.getInt(new ByteArray(values[i])); + } + return indexes; + } + public int[] indexOfMV(Object value) { Object[] multiValues = (Object[]) value; int[] indexes = new int[multiValues.length]; 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/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/readers/DefaultValueColumnReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReader.java index 22a001ae3b..ee70b62013 100644 --- 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 @@ -18,7 +18,6 @@ */ package org.apache.pinot.segment.local.segment.readers; -import java.io.IOException; import javax.annotation.Nullable; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.readers.ColumnReader; @@ -35,6 +34,15 @@ public class DefaultValueColumnReader implements ColumnReader { private final String _columnName; private final int _numDocs; private final Object _defaultValue; + 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; @@ -49,6 +57,7 @@ public DefaultValueColumnReader(String columnName, int numDocs, FieldSpec fieldS _columnName = columnName; _numDocs = numDocs; _currentIndex = 0; + _dataType = fieldSpec.getDataType(); // For multi-value fields, wrap the default value in an array Object defaultNullValue = fieldSpec.getDefaultNullValue(); @@ -56,6 +65,48 @@ public DefaultValueColumnReader(String columnName, int numDocs, FieldSpec fieldS _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; + } } } @@ -66,8 +117,7 @@ public boolean hasNext() { @Override @Nullable - public Object next() - throws IOException { + public Object next() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } @@ -76,19 +126,271 @@ public Object next() } @Override - public void rewind() - throws IOException { + 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 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 int[] nextIntMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + _currentIndex++; + return _defaultIntMV; + } + + @Override + public long[] nextLongMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + _currentIndex++; + return _defaultLongMV; + } + + @Override + public float[] nextFloatMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + _currentIndex++; + return _defaultFloatMV; + } + + @Override + public double[] nextDoubleMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + _currentIndex++; + return _defaultDoubleMV; + } + + @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 void close() - throws IOException { + 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; + } + + // Multi-value accessors + + @Override + public int[] getIntMV(int docId) { + validateDocId(docId); + return _defaultIntMV; + } + + @Override + public long[] getLongMV(int docId) { + validateDocId(docId); + return _defaultLongMV; + } + + @Override + public float[] getFloatMV(int docId) { + validateDocId(docId); + return _defaultFloatMV; + } + + @Override + public double[] getDoubleMV(int docId) { + validateDocId(docId); + return _defaultDoubleMV; + } + + @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 index a69af20d8a..e9d16782fd 100644 --- 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 @@ -22,7 +22,6 @@ 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; @@ -48,7 +47,8 @@ public class PinotSegmentColumnReaderFactory implements ColumnReaderFactory { private final IndexSegment _indexSegment; private Schema _targetSchema; - @Nullable private Map _columnReaders; + private Set _colsToRead; + private Map _columnReaders; /** * Create a PinotSegmentColumnReaderFactory. @@ -63,7 +63,14 @@ public PinotSegmentColumnReaderFactory(IndexSegment indexSegment) { @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()); @@ -130,8 +137,7 @@ public Map getAllColumnReaders() /** * Internal method to initialize all column readers during factory initialization. */ - private Map initializeAllColumnReaders() - throws IOException { + private Map initializeAllColumnReaders() { if (_targetSchema == null) { throw new IllegalStateException("Factory not initialized. Call init() first."); } @@ -139,12 +145,11 @@ private Map initializeAllColumnReaders() Map allReaders = new HashMap<>(); // Create readers for all columns in the target schema - for (FieldSpec fieldSpec : _targetSchema.getAllFieldSpecs()) { + for (String columnName : _colsToRead) { + FieldSpec fieldSpec = _targetSchema.getFieldSpecFor(columnName); if (fieldSpec.isVirtualColumn()) { continue; } - - String columnName = fieldSpec.getName(); ColumnReader reader = createColumnReader(columnName, fieldSpec); allReaders.put(columnName, reader); } 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 index 9a09394466..827839bdee 100644 --- 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 @@ -21,6 +21,7 @@ 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; @@ -38,6 +39,7 @@ 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; @@ -55,6 +57,11 @@ public PinotSegmentColumnReaderImpl(IndexSegment indexSegment, String 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 @@ -86,6 +93,172 @@ public Object next() 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 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; + } + + @Override + public int[] nextIntMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + int[] value = _segmentColumnReader.getIntMV(_currentIndex); + _currentIndex++; + return value; + } + + @Override + public long[] nextLongMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + long[] value = _segmentColumnReader.getLongMV(_currentIndex); + _currentIndex++; + return value; + } + + @Override + public float[] nextFloatMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + float[] value = _segmentColumnReader.getFloatMV(_currentIndex); + _currentIndex++; + return value; + } + + @Override + public double[] nextDoubleMV() { + if (!hasNext()) { + throw new IllegalStateException("No more values available"); + } + double[] value = _segmentColumnReader.getDoubleMV(_currentIndex); + _currentIndex++; + return value; + } + + @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 { @@ -98,6 +271,94 @@ 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); + } + + // Multi-value accessors + + @Override + public int[] getIntMV(int docId) { + return _segmentColumnReader.getIntMV(docId); + } + + @Override + public long[] getLongMV(int docId) { + return _segmentColumnReader.getLongMV(docId); + } + + @Override + public float[] getFloatMV(int docId) { + return _segmentColumnReader.getFloatMV(docId); + } + + @Override + public double[] getDoubleMV(int docId) { + return _segmentColumnReader.getDoubleMV(docId); + } + + @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 { 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 index f65afe106e..82dd918124 100644 --- 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 @@ -77,7 +77,11 @@ public abstract class ColumnarSegmentBuildingTestBase { 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"; @@ -109,7 +113,11 @@ public void setUp() .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(); @@ -125,7 +133,11 @@ public void setUp() .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) @@ -152,20 +164,90 @@ public void tearDown() protected File createRowMajorSegment() throws Exception { - File outputDir = new File(_tempDir, "rowMajorSegment"); + 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, _originalSchema); + SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); config.setOutDir(outputDir.getAbsolutePath()); - config.setSegmentName(SEGMENT_NAME + "_rowMajor"); + 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 + "_rowMajor"); + 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) @@ -423,8 +505,16 @@ protected List generateTestData(int numRows) { // 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); } 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..5078fc6197 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollectorTest.java @@ -0,0 +1,708 @@ +/** + * 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.*; + + +/** + * 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..482b29c543 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/DefaultValueColumnReaderTest.java @@ -0,0 +1,613 @@ +/** + * 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.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()); + int[] result = reader.nextIntMV(); + Assert.assertTrue(Arrays.equals(result, expectedArray)); + } + + // Test random access + reader.rewind(); + for (int i = 0; i < NUM_DOCS; i++) { + int[] result = reader.getIntMV(i); + Assert.assertTrue(Arrays.equals(result, expectedArray)); + } + + // Test that the same array instance is returned (optimization) + reader.rewind(); + int[] firstCall = reader.getIntMV(0); + int[] secondCall = reader.getIntMV(1); + 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++) { + long[] result = reader.nextLongMV(); + Assert.assertTrue(Arrays.equals(result, expectedArray)); + } + + // Test random access and array reuse + reader.rewind(); + long[] firstCall = reader.getLongMV(0); + long[] secondCall = reader.getLongMV(1); + 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++) { + float[] result = reader.nextFloatMV(); + Assert.assertTrue(Arrays.equals(result, expectedArray)); + } + + // Test random access and array reuse + reader.rewind(); + float[] firstCall = reader.getFloatMV(0); + float[] secondCall = reader.getFloatMV(1); + 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++) { + double[] result = reader.nextDoubleMV(); + Assert.assertTrue(Arrays.equals(result, expectedArray)); + } + + // Test random access and array reuse + reader.rewind(); + double[] firstCall = reader.getDoubleMV(0); + double[] secondCall = reader.getDoubleMV(1); + 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/PinotSegmentColumnReaderImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java new file mode 100644 index 0000000000..b2c7c5932a --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReaderImplTest.java @@ -0,0 +1,927 @@ +/** + * 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.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 + Object[][] baseConfigs = new Object[][] { + {MV_INT_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getIntMV, + (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextIntMV, null}, + {MV_LONG_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getLongMV, + (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextLongMV, null}, + {MV_FLOAT_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getFloatMV, + (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextFloatMV, + (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) PinotSegmentColumnReaderImpl::getDoubleMV, + (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextDoubleMV, 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++) { + randomAccessValues[docId] = reader.getIntMV(docId); + } + + // 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/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 index ce3f898a6e..faca8bfac0 100644 --- 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 @@ -29,15 +29,120 @@ * 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 follows an iterator pattern similar to RecordReader: + *

    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() and next()
      • - *
    • Rewind capability to restart iteration
      • - *
    • Resource cleanup
      • + *
    • 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.
      • *
    * - *

    Implementations should handle data type conversions, default values for new columns, - * and efficient column-wise data access patterns. */ public interface ColumnReader extends Closeable, Serializable { @@ -59,13 +164,74 @@ public interface ColumnReader extends Closeable, Serializable { 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 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. + * + * @throws IOException If an I/O error occurs while reading + */ + int[] nextIntMV() throws IOException; + long[] nextLongMV() throws IOException; + float[] nextFloatMV() throws IOException; + double[] 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; + void rewind() throws IOException; /** * Get the name of the column. @@ -73,4 +239,57 @@ void rewind() * @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; + + // 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. + * + * @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[] getIntMV(int docId) throws IOException; + long[] getLongMV(int docId) throws IOException; + float[] getFloatMV(int docId) throws IOException; + double[] 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 index 1df2be955c..eadb1ad656 100644 --- 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 @@ -49,6 +49,15 @@ public interface ColumnReaderFactory extends Closeable, Serializable { 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; + /** * Get the set of column names available in the source data. * From ca5b22ecfe9dafbf8992121474865bf5bb115b79 Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Thu, 18 Dec 2025 09:45:26 +0530 Subject: [PATCH 3/8] Followups on column reader changes (#17293) * Followups on column reader changes * Handle nulls in default columns more appropriately * Add isSingleValue() to ColumnReader * Minor refactoring in EpochTimeHandler * Avoid null changes of DefaultValueColumnReader in this PR (cherry picked from commit a9f36e438cec8ab80139acf03993e39947582c66) --- .../timehandler/EpochTimeHandler.java | 44 +++++++++++++ .../readers/DefaultValueColumnReader.java | 13 ++++ .../PinotSegmentColumnReaderFactory.java | 9 +-- .../readers/PinotSegmentColumnReaderImpl.java | 11 ++++ .../PinotSegmentColumnarDataSource.java | 61 +++++++++++++++++ .../pinot/spi/data/readers/ColumnReader.java | 22 ++++++- .../spi/data/readers/ColumnReaderFactory.java | 6 +- .../spi/data/readers/ColumnarDataSource.java | 66 +++++++++++++++++++ 8 files changed, 221 insertions(+), 11 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnarDataSource.java create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/ColumnarDataSource.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java index 9e6ad822b5..a3fbd47b60 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java @@ -70,4 +70,48 @@ public String handleTime(GenericRow row) { return DEFAULT_PARTITION; } } + + @Override + public String getTimeColumn() { + return _timeColumn; + } + + @Override + @Nullable + public String handleTimeColumn(Object columnValue) { + long timeMs = _formatSpec.fromFormatToMillis(columnValue.toString()); + + // Apply time filter + if (_startTimeMs > 0) { + boolean outsideTimeWindow = (timeMs < _startTimeMs || timeMs >= _endTimeMs); + if (outsideTimeWindow != _negateWindowFilter) { + return null; + } + } + + // Round time if needed + if (_roundBucketMs > 0) { + timeMs = (timeMs / _roundBucketMs) * _roundBucketMs; + } + + // Compute partition + if (_partitionBucketMs > 0) { + return Long.toString(timeMs / _partitionBucketMs); + } else { + return DEFAULT_PARTITION; + } + } + + @Override + @Nullable + public Object getModifiedTimeValue(Object columnValue) { + // Round time if needed + if (_roundBucketMs > 0) { + long timeMs = _formatSpec.fromFormatToMillis(columnValue.toString()); + timeMs = (timeMs / _roundBucketMs) * _roundBucketMs; + return _dataType.convert(_formatSpec.fromMillisToFormat(timeMs)); + } else { + return columnValue; + } + } } 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 index ee70b62013..6e97958f96 100644 --- 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 @@ -34,6 +34,7 @@ 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 @@ -57,6 +58,7 @@ public DefaultValueColumnReader(String columnName, int numDocs, FieldSpec fieldS _columnName = columnName; _numDocs = numDocs; _currentIndex = 0; + _fieldSpec = fieldSpec; _dataType = fieldSpec.getDataType(); // For multi-value fields, wrap the default value in an array @@ -146,6 +148,11 @@ public void skipNext() { _currentIndex++; } + @Override + public boolean isSingleValue() { + return _fieldSpec.isSingleValueField(); + } + @Override public boolean isInt() { return _dataType == FieldSpec.DataType.INT; @@ -338,6 +345,12 @@ public byte[] getBytes(int docId) { return (byte[]) _defaultValue; } + @Override + public Object getValue(int docId) { + validateDocId(docId); + return _defaultValue; + } + // Multi-value accessors @Override 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 index e9d16782fd..2b40444a52 100644 --- 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 @@ -91,17 +91,12 @@ public int getNumDocs() { } @Override - public ColumnReader getColumnReader(String columnName) - throws IOException { + public ColumnReader getColumnReader(String columnName) { if (_targetSchema == null) { throw new IllegalStateException("Factory not initialized. Call init() first."); } - ColumnReader reader = _columnReaders.get(columnName); - if (reader == null) { - throw new IOException("Column reader not found for column: " + columnName); - } - return reader; + return _columnReaders.get(columnName); } /** 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 index 827839bdee..2af71982e5 100644 --- 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 @@ -109,6 +109,11 @@ public void skipNext() { _currentIndex++; } + @Override + public boolean isSingleValue() { + return _segmentColumnReader.isSingleValue(); + } + @Override public boolean isInt() { return _dataType == FieldSpec.DataType.INT; @@ -314,6 +319,12 @@ 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 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..8b3306f694 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnarDataSource.java @@ -0,0 +1,61 @@ +/** + * 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; + + public PinotSegmentColumnarDataSource(IndexSegment indexSegment) { + _indexSegment = indexSegment; + _totalDocs = indexSegment.getSegmentMetadata().getTotalDocs(); + } + + @Override + public int getTotalDocs() { + return _totalDocs; + } + + @Override + public ColumnReaderFactory createColumnReaderFactory() { + return new PinotSegmentColumnReaderFactory(_indexSegment); + } + + @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-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 index faca8bfac0..28bc576acb 100644 --- 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 @@ -188,6 +188,11 @@ Object next() */ 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. @@ -262,7 +267,7 @@ Object next() /** * 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. + * 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 @@ -275,6 +280,21 @@ Object next() 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 /** 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 index eadb1ad656..e31293c0b9 100644 --- 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 @@ -23,6 +23,7 @@ import java.io.Serializable; import java.util.Map; import java.util.Set; +import javax.annotation.Nullable; import org.apache.pinot.spi.data.Schema; @@ -70,11 +71,10 @@ void init(Schema targetSchema, Set colsToRead) * 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) - * @throws IOException If the column reader cannot be obtained */ - ColumnReader getColumnReader(String columnName) - throws IOException; + @Nullable ColumnReader getColumnReader(String columnName); /** * Get all column readers for the target schema. 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(); +} From fdef66a3b05e73785adad83c1eb7ae9787116b14 Mon Sep 17 00:00:00 2001 From: 9aman <35227405+9aman@users.noreply.github.com> Date: Fri, 19 Dec 2025 17:34:58 +0530 Subject: [PATCH 4/8] Fixing ColumnReader interface to return validity bitset in case of multi-value primitive type (#17387) * Fixing ColumnReader interface to return validity bitset in case of multi-value primitive type * Minor formatting improvement * 1. Improve test cases to assert that nulti-value columns have no nulls. 2. Add documentation stating the reason for passing null BitSet validity for multi-value primitive columns. * Renaming validity to nulls to capture the intent of the variable better (cherry picked from commit bdd834797a7cef5f647b84b9fdf395f3e2678475) --- .../readers/DefaultValueColumnReader.java | 33 ++--- .../readers/PinotSegmentColumnReaderImpl.java | 36 ++--- .../readers/DefaultValueColumnReaderTest.java | 42 +++--- .../PinotSegmentColumnReaderImplTest.java | 54 ++++++-- .../pinot/spi/data/readers/ColumnReader.java | 61 +++++++-- .../spi/data/readers/MultiValueResult.java | 123 ++++++++++++++++++ 6 files changed, 279 insertions(+), 70 deletions(-) create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/data/readers/MultiValueResult.java 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 index 6e97958f96..4392fe9cd0 100644 --- 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 @@ -21,6 +21,7 @@ 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; /** @@ -238,39 +239,39 @@ public byte[] nextBytes() { } @Override - public int[] nextIntMV() { + public MultiValueResult nextIntMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } _currentIndex++; - return _defaultIntMV; + return MultiValueResult.of(_defaultIntMV, null); } @Override - public long[] nextLongMV() { + public MultiValueResult nextLongMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } _currentIndex++; - return _defaultLongMV; + return MultiValueResult.of(_defaultLongMV, null); } @Override - public float[] nextFloatMV() { + public MultiValueResult nextFloatMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } _currentIndex++; - return _defaultFloatMV; + return MultiValueResult.of(_defaultFloatMV, null); } @Override - public double[] nextDoubleMV() { + public MultiValueResult nextDoubleMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } _currentIndex++; - return _defaultDoubleMV; + return MultiValueResult.of(_defaultDoubleMV, null); } @Override @@ -354,27 +355,27 @@ public Object getValue(int docId) { // Multi-value accessors @Override - public int[] getIntMV(int docId) { + public MultiValueResult getIntMV(int docId) { validateDocId(docId); - return _defaultIntMV; + return MultiValueResult.of(_defaultIntMV, null); } @Override - public long[] getLongMV(int docId) { + public MultiValueResult getLongMV(int docId) { validateDocId(docId); - return _defaultLongMV; + return MultiValueResult.of(_defaultLongMV, null); } @Override - public float[] getFloatMV(int docId) { + public MultiValueResult getFloatMV(int docId) { validateDocId(docId); - return _defaultFloatMV; + return MultiValueResult.of(_defaultFloatMV, null); } @Override - public double[] getDoubleMV(int docId) { + public MultiValueResult getDoubleMV(int docId) { validateDocId(docId); - return _defaultDoubleMV; + return MultiValueResult.of(_defaultDoubleMV, null); } @Override 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 index 2af71982e5..b34a0e53ef 100644 --- 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 @@ -23,6 +23,7 @@ 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; /** @@ -204,44 +205,47 @@ public byte[] nextBytes() { 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 int[] nextIntMV() { + public MultiValueResult nextIntMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } int[] value = _segmentColumnReader.getIntMV(_currentIndex); _currentIndex++; - return value; + return MultiValueResult.of(value, null); } @Override - public long[] nextLongMV() { + public MultiValueResult nextLongMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } long[] value = _segmentColumnReader.getLongMV(_currentIndex); _currentIndex++; - return value; + return MultiValueResult.of(value, null); } @Override - public float[] nextFloatMV() { + public MultiValueResult nextFloatMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } float[] value = _segmentColumnReader.getFloatMV(_currentIndex); _currentIndex++; - return value; + return MultiValueResult.of(value, null); } @Override - public double[] nextDoubleMV() { + public MultiValueResult nextDoubleMV() { if (!hasNext()) { throw new IllegalStateException("No more values available"); } double[] value = _segmentColumnReader.getDoubleMV(_currentIndex); _currentIndex++; - return value; + return MultiValueResult.of(value, null); } @Override @@ -328,23 +332,23 @@ public Object getValue(int docId) // Multi-value accessors @Override - public int[] getIntMV(int docId) { - return _segmentColumnReader.getIntMV(docId); + public MultiValueResult getIntMV(int docId) { + return MultiValueResult.of(_segmentColumnReader.getIntMV(docId), null); } @Override - public long[] getLongMV(int docId) { - return _segmentColumnReader.getLongMV(docId); + public MultiValueResult getLongMV(int docId) { + return MultiValueResult.of(_segmentColumnReader.getLongMV(docId), null); } @Override - public float[] getFloatMV(int docId) { - return _segmentColumnReader.getFloatMV(docId); + public MultiValueResult getFloatMV(int docId) { + return MultiValueResult.of(_segmentColumnReader.getFloatMV(docId), null); } @Override - public double[] getDoubleMV(int docId) { - return _segmentColumnReader.getDoubleMV(docId); + public MultiValueResult getDoubleMV(int docId) { + return MultiValueResult.of(_segmentColumnReader.getDoubleMV(docId), null); } @Override 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 index 482b29c543..abe05a48af 100644 --- 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 @@ -23,6 +23,7 @@ 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; @@ -222,21 +223,23 @@ public void testMultiValueIntColumn() throws IOException { int[] expectedArray = new int[]{expectedValue}; for (int i = 0; i < NUM_DOCS; i++) { Assert.assertTrue(reader.hasNext()); - int[] result = reader.nextIntMV(); - Assert.assertTrue(Arrays.equals(result, expectedArray)); + 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++) { - int[] result = reader.getIntMV(i); - Assert.assertTrue(Arrays.equals(result, expectedArray)); + 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); - int[] secondCall = reader.getIntMV(1); + int[] firstCall = reader.getIntMV(0).getValues(); + int[] secondCall = reader.getIntMV(1).getValues(); Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused"); reader.close(); @@ -254,14 +257,15 @@ public void testMultiValueLongColumn() throws IOException { long expectedValue = ((Number) fieldSpec.getDefaultNullValue()).longValue(); long[] expectedArray = new long[]{expectedValue}; for (int i = 0; i < NUM_DOCS; i++) { - long[] result = reader.nextLongMV(); - Assert.assertTrue(Arrays.equals(result, expectedArray)); + 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); - long[] secondCall = reader.getLongMV(1); + 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"); @@ -277,14 +281,15 @@ public void testMultiValueFloatColumn() throws IOException { float expectedValue = ((Number) fieldSpec.getDefaultNullValue()).floatValue(); float[] expectedArray = new float[]{expectedValue}; for (int i = 0; i < NUM_DOCS; i++) { - float[] result = reader.nextFloatMV(); - Assert.assertTrue(Arrays.equals(result, expectedArray)); + 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); - float[] secondCall = reader.getFloatMV(1); + float[] firstCall = reader.getFloatMV(0).getValues(); + float[] secondCall = reader.getFloatMV(1).getValues(); Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused"); reader.close(); @@ -299,14 +304,15 @@ public void testMultiValueDoubleColumn() throws IOException { double expectedValue = ((Number) fieldSpec.getDefaultNullValue()).doubleValue(); double[] expectedArray = new double[]{expectedValue}; for (int i = 0; i < NUM_DOCS; i++) { - double[] result = reader.nextDoubleMV(); - Assert.assertTrue(Arrays.equals(result, expectedArray)); + 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); - double[] secondCall = reader.getDoubleMV(1); + double[] firstCall = reader.getDoubleMV(0).getValues(); + double[] secondCall = reader.getDoubleMV(1).getValues(); Assert.assertSame(firstCall, secondCall, "Multi-value arrays should be reused"); reader.close(); 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 index b2c7c5932a..f3e8d0141f 100644 --- 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 @@ -28,6 +28,7 @@ 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; @@ -359,13 +360,39 @@ public void testSingleValueAccessors(String columnName, SingleValueGetter random @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) PinotSegmentColumnReaderImpl::getIntMV, - (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextIntMV, null}, - {MV_LONG_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getLongMV, - (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextLongMV, null}, - {MV_FLOAT_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getFloatMV, - (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextFloatMV, + {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++) { @@ -373,8 +400,16 @@ public Object[][] multiValueAccessorProvider() { } return expectedFloatArray; } }, - {MV_DOUBLE_COL, (MultiValueGetter) PinotSegmentColumnReaderImpl::getDoubleMV, - (MultiValueSequentialGetter) PinotSegmentColumnReaderImpl::nextDoubleMV, null}, + {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, @@ -647,7 +682,8 @@ public void testIteratorAndRandomAccessConsistencyMV() // First, read all values using random access int[][] randomAccessValues = new int[totalDocs][]; for (int docId = 0; docId < totalDocs; docId++) { - randomAccessValues[docId] = reader.getIntMV(docId); + MultiValueResult result = reader.getIntMV(docId); + randomAccessValues[docId] = result.getValues(); } // Now read using iterator pattern 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 index 28bc576acb..3c05c26eba 100644 --- 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 @@ -57,7 +57,7 @@ * } catch (Exception e) { * // Handle exception / log * continue; -* } + * } * if (value != null) { * // Process non-null value * processValue(value); @@ -200,10 +200,15 @@ Object next() * Otherwise, clients should use next() and cast the result. */ boolean isInt(); + boolean isLong(); + boolean isFloat(); + boolean isDouble(); + boolean isString(); + boolean isBytes(); /** @@ -212,24 +217,44 @@ Object next() * @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 */ - int[] nextIntMV() throws IOException; - long[] nextLongMV() throws IOException; - float[] nextFloatMV() throws IOException; - double[] nextDoubleMV() throws IOException; - String[] nextStringMV() throws IOException; - byte[][] nextBytesMV() throws IOException; + 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. @@ -274,10 +299,15 @@ Object next() * @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; /** @@ -302,14 +332,23 @@ Object next() * 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 */ - int[] getIntMV(int docId) throws IOException; - long[] getLongMV(int docId) throws IOException; - float[] getFloatMV(int docId) throws IOException; - double[] getDoubleMV(int docId) throws IOException; + 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/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; + } +} From 0825eccc250d0b291bb8fa081143b21eff74ffab Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Mon, 29 Dec 2025 14:39:40 +0530 Subject: [PATCH 5/8] Add InMemoryColumnReader for testing and update PinotSegmentColumnReaderFactory to support configurable default value readers (#17415) * Add InMemoryColumnReader for testing and update PinotSegmentColumnReaderFactory to support configurable default value readers * Remove default in PinotSegmentColumnarDataSource (cherry picked from commit 9f248c52aff01ff6740dfbab5a72d6b185ba1e2c) --- .../PinotSegmentColumnReaderFactory.java | 31 ++- .../PinotSegmentColumnarDataSource.java | 6 +- .../segment/readers/InMemoryColumnReader.java | 258 ++++++++++++++++++ 3 files changed, 288 insertions(+), 7 deletions(-) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/InMemoryColumnReader.java 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 index 2b40444a52..b6eb206c31 100644 --- 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 @@ -22,6 +22,7 @@ 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; @@ -49,6 +50,7 @@ public class PinotSegmentColumnReaderFactory implements ColumnReaderFactory { private Schema _targetSchema; private Set _colsToRead; private Map _columnReaders; + private final boolean _initializeDefaultValueReaders; /** * Create a PinotSegmentColumnReaderFactory. @@ -56,8 +58,22 @@ public class PinotSegmentColumnReaderFactory implements ColumnReaderFactory { * @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 @@ -103,21 +119,24 @@ public ColumnReader getColumnReader(String 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; + 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 { - // New column - create a default value reader - LOGGER.debug("Creating default value reader for new column: {}", columnName); - columnReader = new DefaultValueColumnReader(columnName, getNumDocs(), targetFieldSpec); + 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; @@ -146,7 +165,9 @@ private Map initializeAllColumnReaders() { continue; } ColumnReader reader = createColumnReader(columnName, fieldSpec); - allReaders.put(columnName, reader); + if (reader != null) { + allReaders.put(columnName, reader); + } } return allReaders; 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 index 8b3306f694..960dc7ebc4 100644 --- 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 @@ -32,10 +32,12 @@ public class PinotSegmentColumnarDataSource implements ColumnarDataSource { private final IndexSegment _indexSegment; private final int _totalDocs; + private final boolean _initializeDefaultValueReaders; - public PinotSegmentColumnarDataSource(IndexSegment indexSegment) { + public PinotSegmentColumnarDataSource(IndexSegment indexSegment, boolean initializeDefaultValueReaders) { _indexSegment = indexSegment; _totalDocs = indexSegment.getSegmentMetadata().getTotalDocs(); + _initializeDefaultValueReaders = initializeDefaultValueReaders; } @Override @@ -45,7 +47,7 @@ public int getTotalDocs() { @Override public ColumnReaderFactory createColumnReaderFactory() { - return new PinotSegmentColumnReaderFactory(_indexSegment); + return new PinotSegmentColumnReaderFactory(_indexSegment, _initializeDefaultValueReaders); } @Override 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 + } +} From 32eeff8793c84188cb5eb49204c2af48f69c2e93 Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Thu, 5 Feb 2026 19:53:14 +0530 Subject: [PATCH 6/8] Add configs when initing column reader factory (#17635) (cherry picked from commit 783075165da49ee23b8e5e56215ddf9697926b35) --- .../spi/data/readers/ColumnReaderFactory.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 index e31293c0b9..ce9a9c3cc9 100644 --- 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 @@ -41,6 +41,12 @@ */ 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. * @@ -59,6 +65,18 @@ void init(Schema targetSchema) 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. * From 202aa2ee6b38b0974f85a3804ed0c9eaaf3d4bb6 Mon Sep 17 00:00:00 2001 From: Tony Song Date: Wed, 27 May 2026 20:32:05 -0700 Subject: [PATCH 7/8] Replace star import in AbstractColumnStatisticsCollectorTest Apache uses `import static org.testng.Assert.*;` in this test file which violates linkedin/pinot's AvoidStarImport checkstyle rule. Replace with explicit imports for the methods actually used. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../impl/stats/AbstractColumnStatisticsCollectorTest.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 index 5078fc6197..3be1c48aaf 100644 --- 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 @@ -33,7 +33,11 @@ import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import static org.testng.Assert.*; +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; /** From 15e03d6f94480d7bd0d71ab41e398a2675262e46 Mon Sep 17 00:00:00 2001 From: Tony Song Date: Wed, 27 May 2026 20:50:20 -0700 Subject: [PATCH 8/8] Fix CI compile error: remove orphan TimeHandler methods + typo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The EpochTimeHandler additions from cherry-pick apache/pinot#17293 included getTimeColumn(), handleTimeColumn(), getModifiedTimeValue() with @Override annotations. These methods depend on apache/pinot#17171 which extended the TimeHandler interface — that PR is not in linkedin/pinot, so the @Override annotations don't override anything and the build fails. Remove the orphan methods entirely. apache/pinot#17293's actual intended change to EpochTimeHandler was a lazy-init refactor of getModifiedTimeValue, which is moot here since the method doesn't exist in linkedin/pinot's TimeHandler contract. Also: fix "recort" → "record" typo in SegmentIndexCreationDriverImpl javadoc inherited from apache/pinot#16727 (still present in apache HEAD). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../timehandler/EpochTimeHandler.java | 44 ------------------- .../impl/SegmentIndexCreationDriverImpl.java | 2 +- 2 files changed, 1 insertion(+), 45 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java index a3fbd47b60..9e6ad822b5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/timehandler/EpochTimeHandler.java @@ -70,48 +70,4 @@ public String handleTime(GenericRow row) { return DEFAULT_PARTITION; } } - - @Override - public String getTimeColumn() { - return _timeColumn; - } - - @Override - @Nullable - public String handleTimeColumn(Object columnValue) { - long timeMs = _formatSpec.fromFormatToMillis(columnValue.toString()); - - // Apply time filter - if (_startTimeMs > 0) { - boolean outsideTimeWindow = (timeMs < _startTimeMs || timeMs >= _endTimeMs); - if (outsideTimeWindow != _negateWindowFilter) { - return null; - } - } - - // Round time if needed - if (_roundBucketMs > 0) { - timeMs = (timeMs / _roundBucketMs) * _roundBucketMs; - } - - // Compute partition - if (_partitionBucketMs > 0) { - return Long.toString(timeMs / _partitionBucketMs); - } else { - return DEFAULT_PARTITION; - } - } - - @Override - @Nullable - public Object getModifiedTimeValue(Object columnValue) { - // Round time if needed - if (_roundBucketMs > 0) { - long timeMs = _formatSpec.fromFormatToMillis(columnValue.toString()); - timeMs = (timeMs / _roundBucketMs) * _roundBucketMs; - return _dataType.convert(_formatSpec.fromMillisToFormat(timeMs)); - } else { - return columnValue; - } - } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java index a7c8857a7c..e33f1570d5 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java @@ -639,7 +639,7 @@ public SegmentPreIndexStatsContainer getSegmentStats() { * Build segment using columnar approach. * This method builds the segment by processing data column-wise instead of row-wise. * Following is not supported: - *

  • recort transformation + *
  • record transformation *
  • sorted column change wrt to input data * *

    Initialize the driver using {@link #init(SegmentGeneratorConfig, ColumnReaderFactory)}