Cherry-pick: columnar segment build SPI from apache/pinot (6 PRs)

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.
+ *
+ * <p>This data source enables columnar segment building by providing access to column readers
+ * instead of row-based record readers. It supports:
+ * <ul>
+ *   <li>Columnar statistics collection</li>
+ *   <li>Column-wise data access through ColumnReaderFactory</li>
+ *   <li>Efficient handling of new columns with default values</li>
+ *   <li>Data type conversions during schema evolution</li>
+ * </ul>
+ */
+public class ColumnarSegmentCreationDataSource implements SegmentCreationDataSource, Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ColumnarSegmentCreationDataSource.class);
+
+  private final Map<String, ColumnReader> _columnReaders;
+
+  /**
+   * Create a ColumnarSegmentCreationDataSource.
+   *
+   * @param columnReaders Map of column name to ColumnReader instances
+   */
+  public ColumnarSegmentCreationDataSource(Map<String, ColumnReader> 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<String, ColumnReader> getColumnReaders() {
+    return _columnReaders;
+  }
+
+  @Override
+  public void close()
+      throws IOException {
+    for (ColumnReader columnReader : _columnReaders.values()) {
+      columnReader.close();
+    }
+  }
+}

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<IndexType<?, ?, ?>, 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<IndexType<?, ?, ?>, IndexCreator> creatorsByIndex, String columnName, FieldSpec fieldSpec,
       SegmentDictionaryCreator dictionaryCreator, int sourceDocId, int onDiskDocPos,

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];