real-mj-song commented on code in PR #18638: URL: https://github.com/apache/pinot/pull/18638#discussion_r3408627570
########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowAccumulators.java: ########## @@ -0,0 +1,184 @@ +/** + * 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.plugin.inputformat.arrow; + +import com.google.common.base.Preconditions; +import java.io.IOException; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.util.VectorAppender; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.ColumnReader; + + +/** + * Package-private helper shared by {@link ArrowColumnReaderFactory} and + * {@link InMemoryArrowColumnReaderFactory}: walks every record batch in an {@link ArrowReader}, + * concatenates each wanted column's values into a per-column accumulator {@link FieldVector} via + * Arrow's {@link VectorAppender}, and produces one {@link ArrowColumnReader} per accumulator. Review Comment: Fixed — the `@link` now points at `ArrowFileColumnReaderFactory`. The helper is shared between the caller-managed `ArrowColumnReaderFactory` and the file-backed `ArrowFileColumnReaderFactory`. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowAccumulators.java: ########## @@ -0,0 +1,184 @@ +/** + * 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.plugin.inputformat.arrow; + +import com.google.common.base.Preconditions; +import java.io.IOException; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.util.VectorAppender; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.ColumnReader; + + +/** + * Package-private helper shared by {@link ArrowColumnReaderFactory} and + * {@link InMemoryArrowColumnReaderFactory}: walks every record batch in an {@link ArrowReader}, + * concatenates each wanted column's values into a per-column accumulator {@link FieldVector} via + * Arrow's {@link VectorAppender}, and produces one {@link ArrowColumnReader} per accumulator. + * + * <p>Accumulator vectors are allocated against the caller-supplied {@link BufferAllocator}. The + * caller (factory) owns and closes them via {@link Result#getAccumulators()}; this helper does + * not retain references. + */ +final class ArrowAccumulators { + + private ArrowAccumulators() { + } + + static Result populate(ArrowReader reader, BufferAllocator allocator, Schema targetSchema, + @Nullable Set<String> colsToRead) + throws IOException { + Set<String> wantedColumns = computeWantedColumns(targetSchema, colsToRead); + + VectorSchemaRoot perBatchRoot = reader.getVectorSchemaRoot(); + Set<String> availableColumns = Collections.unmodifiableSet(collectAvailableNames(perBatchRoot)); + + Map<String, FieldVector> accumulators = new LinkedHashMap<>(); + Map<String, VectorAppender> appenders = new LinkedHashMap<>(); + for (FieldVector source : perBatchRoot.getFieldVectors()) { + String name = source.getField().getName(); + if (!wantedColumns.isEmpty() && !wantedColumns.contains(name)) { + continue; + } + // Dictionary-encoded columns surface as their index type (e.g. Int32) rather than the + // decoded logical type. Reject loudly so we don't silently produce a wrong segment. The + // row-major ArrowRecordExtractor decodes via DictionaryEncoder.decode; adding the same + // here is left as a follow-up once a real use case appears. + Preconditions.checkArgument(source.getField().getDictionary() == null, + "Dictionary-encoded Arrow column '%s' is not supported by Arrow column-major build. " + + "Use ArrowRecordReader (row-major) for files containing dictionary-encoded columns.", + name); + FieldVector accumulator = source.getField().createVector(allocator); + // Pre-allocate buffers so VectorAppender can read offsets / validity from them on the + // first append (otherwise BaseVariableWidthVector visits IOOBE on an empty offset buffer). + accumulator.allocateNew(); + accumulators.put(name, accumulator); + appenders.put(name, new VectorAppender(accumulator)); + } + + // Walk every record batch and bulk-append each wanted column into its accumulator via + // Arrow's VectorAppender (Visitor-based; grows offset / data buffers once per batch and + // bulk-copies, rather than per-row copyValueSafe). Single-batch and multi-batch inputs go + // through the same code path — VectorAppender handles either correctly. + boolean anyBatchSeen = false; + while (reader.loadNextBatch()) { + if (perBatchRoot.getRowCount() == 0) { + continue; + } + anyBatchSeen = true; + for (FieldVector source : perBatchRoot.getFieldVectors()) { + VectorAppender appender = appenders.get(source.getField().getName()); + if (appender != null) { + source.accept(appender, null); + } + } + } + if (!anyBatchSeen) { + throw new IOException("Arrow source contains no non-empty record batches"); + } Review Comment: Fixed — `populate()` now wraps the accumulator allocation and the per-batch walk in a try/catch; on any throw (missing dictionary, `loadNextBatch()` failure, appender error) `closeAll(accumulators)` releases the already-allocated off-heap vectors before the exception propagates, with the close failure added as suppressed — so the caller's allocator close no longer trips on outstanding buffers. The `!anyBatchSeen` throw was also removed: an empty source now yields zero-doc readers, matching the row-major path. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowColumnReader.java: ########## @@ -0,0 +1,595 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.util.BitSet; +import javax.annotation.Nullable; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; + + +/** + * Column reader for Apache Arrow {@link FieldVector}. + * + * <p>Wraps a single Arrow {@link FieldVector} and exposes sequential and random-access + * read patterns conforming to the {@link ColumnReader} contract. The vector is owned by + * the enclosing {@link ArrowColumnReaderFactory} which is responsible for its lifecycle; + * closing this reader is a no-op for the underlying vector. + * + * <p>Supported Arrow types map to Pinot's stored types as follows: + * <ul> + * <li>{@code IntVector} / {@code BitVector} → {@code INT} (BitVector promoted to 0/1)</li> + * <li>{@code BigIntVector} → {@code LONG}</li> + * <li>{@code Float4Vector} → {@code FLOAT}</li> + * <li>{@code Float8Vector} → {@code DOUBLE}</li> + * <li>{@code DecimalVector} → {@code BIG_DECIMAL}</li> + * <li>{@code VarCharVector} → {@code STRING}</li> + * <li>{@code VarBinaryVector} → {@code BYTES}</li> + * <li>{@code ListVector} of the above → multi-value variant</li> + * </ul> + * + * <p>The list above applies to the typed primitive accessors only ({@link #getInt}, + * {@link #getString}, {@link #getIntMV}, ...). Complex types (Map, Struct, Union, ...) + * and temporal types are still readable via the generic {@link #getValue(int)} / + * {@link #next()} accessors, which delegate to {@link ArrowToPinotTypeConverter} and + * return the same canonical JDK types as the row-major path ({@link + * ArrowRecordExtractor}) — e.g. {@code Map<String, Object>} for Struct / Map, + * {@code Object[]} for List variants, {@code LocalDate} / {@code LocalTime} / + * {@code Timestamp} for temporal types. + * + * <p>This class is not thread-safe. + */ +public class ArrowColumnReader implements ColumnReader { + + private final String _columnName; + private final FieldVector _vector; + private final int _totalDocs; + private final boolean _isSingleValue; + + private int _nextDocId; + + /** + * Construct an ArrowColumnReader for the given vector. + * + * @param columnName Pinot column name + * @param vector Arrow field vector backing this column + */ + public ArrowColumnReader(String columnName, FieldVector vector) { + _columnName = columnName; + _vector = vector; + _totalDocs = vector.getValueCount(); + _isSingleValue = !(vector instanceof ListVector); + _nextDocId = 0; + } + + @Override + public boolean hasNext() { + return _nextDocId < _totalDocs; + } + + @Override + @Nullable + public Object next() + throws IOException { + Object value = getValue(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public boolean isNextNull() + throws IOException { + return _vector.isNull(_nextDocId); + } + + @Override + public void skipNext() + throws IOException { + _nextDocId++; + } Review Comment: Fixed — added a `requireHasNext()` helper invoked by every sequential accessor (`next`, the typed `nextXxx`, `isNextNull`, `skipNext`) so they throw `IllegalStateException("No more values available")` past the end, matching `PinotSegmentColumnReaderImpl` / `DefaultValueColumnReader`. Covered by `testSequentialAccessorsThrowPastEnd`. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowColumnReader.java: ########## @@ -0,0 +1,595 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.util.BitSet; +import javax.annotation.Nullable; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; + + +/** + * Column reader for Apache Arrow {@link FieldVector}. + * + * <p>Wraps a single Arrow {@link FieldVector} and exposes sequential and random-access + * read patterns conforming to the {@link ColumnReader} contract. The vector is owned by + * the enclosing {@link ArrowColumnReaderFactory} which is responsible for its lifecycle; + * closing this reader is a no-op for the underlying vector. + * + * <p>Supported Arrow types map to Pinot's stored types as follows: + * <ul> + * <li>{@code IntVector} / {@code BitVector} → {@code INT} (BitVector promoted to 0/1)</li> + * <li>{@code BigIntVector} → {@code LONG}</li> + * <li>{@code Float4Vector} → {@code FLOAT}</li> + * <li>{@code Float8Vector} → {@code DOUBLE}</li> + * <li>{@code DecimalVector} → {@code BIG_DECIMAL}</li> + * <li>{@code VarCharVector} → {@code STRING}</li> + * <li>{@code VarBinaryVector} → {@code BYTES}</li> + * <li>{@code ListVector} of the above → multi-value variant</li> + * </ul> + * + * <p>The list above applies to the typed primitive accessors only ({@link #getInt}, + * {@link #getString}, {@link #getIntMV}, ...). Complex types (Map, Struct, Union, ...) + * and temporal types are still readable via the generic {@link #getValue(int)} / + * {@link #next()} accessors, which delegate to {@link ArrowToPinotTypeConverter} and + * return the same canonical JDK types as the row-major path ({@link + * ArrowRecordExtractor}) — e.g. {@code Map<String, Object>} for Struct / Map, + * {@code Object[]} for List variants, {@code LocalDate} / {@code LocalTime} / + * {@code Timestamp} for temporal types. + * + * <p>This class is not thread-safe. + */ +public class ArrowColumnReader implements ColumnReader { + + private final String _columnName; + private final FieldVector _vector; + private final int _totalDocs; + private final boolean _isSingleValue; + + private int _nextDocId; + + /** + * Construct an ArrowColumnReader for the given vector. + * + * @param columnName Pinot column name + * @param vector Arrow field vector backing this column + */ + public ArrowColumnReader(String columnName, FieldVector vector) { + _columnName = columnName; + _vector = vector; + _totalDocs = vector.getValueCount(); + _isSingleValue = !(vector instanceof ListVector); + _nextDocId = 0; + } + + @Override + public boolean hasNext() { + return _nextDocId < _totalDocs; + } + + @Override + @Nullable + public Object next() + throws IOException { + Object value = getValue(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public boolean isNextNull() + throws IOException { + return _vector.isNull(_nextDocId); + } + + @Override + public void skipNext() + throws IOException { + _nextDocId++; + } + + @Override + public boolean isSingleValue() { + return _isSingleValue; + } + + @Override + public boolean isInt() { + return _isSingleValue && (_vector instanceof IntVector || _vector instanceof BitVector); + } + + @Override + public boolean isLong() { + return _isSingleValue && _vector instanceof BigIntVector; + } + + @Override + public boolean isFloat() { + return _isSingleValue && _vector instanceof Float4Vector; + } + + @Override + public boolean isDouble() { + return _isSingleValue && _vector instanceof Float8Vector; + } + + @Override + public boolean isBigDecimal() { + return _isSingleValue && _vector instanceof DecimalVector; + } Review Comment: Fixed — these predicates now inspect the `ListVector` element vector for multi-value columns rather than being gated on `_isSingleValue`, so an INT multi-value column reports `isInt()`, a LONG one `isLong()`, etc. New test `testMultiValueColumnTypePredicates` asserts the element-type predicates on an MV reader. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowColumnReader.java: ########## @@ -0,0 +1,595 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.util.BitSet; +import javax.annotation.Nullable; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; + + +/** + * Column reader for Apache Arrow {@link FieldVector}. + * + * <p>Wraps a single Arrow {@link FieldVector} and exposes sequential and random-access + * read patterns conforming to the {@link ColumnReader} contract. The vector is owned by + * the enclosing {@link ArrowColumnReaderFactory} which is responsible for its lifecycle; + * closing this reader is a no-op for the underlying vector. + * + * <p>Supported Arrow types map to Pinot's stored types as follows: + * <ul> + * <li>{@code IntVector} / {@code BitVector} → {@code INT} (BitVector promoted to 0/1)</li> + * <li>{@code BigIntVector} → {@code LONG}</li> + * <li>{@code Float4Vector} → {@code FLOAT}</li> + * <li>{@code Float8Vector} → {@code DOUBLE}</li> + * <li>{@code DecimalVector} → {@code BIG_DECIMAL}</li> + * <li>{@code VarCharVector} → {@code STRING}</li> + * <li>{@code VarBinaryVector} → {@code BYTES}</li> + * <li>{@code ListVector} of the above → multi-value variant</li> + * </ul> + * + * <p>The list above applies to the typed primitive accessors only ({@link #getInt}, + * {@link #getString}, {@link #getIntMV}, ...). Complex types (Map, Struct, Union, ...) + * and temporal types are still readable via the generic {@link #getValue(int)} / + * {@link #next()} accessors, which delegate to {@link ArrowToPinotTypeConverter} and + * return the same canonical JDK types as the row-major path ({@link + * ArrowRecordExtractor}) — e.g. {@code Map<String, Object>} for Struct / Map, + * {@code Object[]} for List variants, {@code LocalDate} / {@code LocalTime} / + * {@code Timestamp} for temporal types. + * + * <p>This class is not thread-safe. + */ +public class ArrowColumnReader implements ColumnReader { + + private final String _columnName; + private final FieldVector _vector; + private final int _totalDocs; + private final boolean _isSingleValue; + + private int _nextDocId; + + /** + * Construct an ArrowColumnReader for the given vector. + * + * @param columnName Pinot column name + * @param vector Arrow field vector backing this column + */ + public ArrowColumnReader(String columnName, FieldVector vector) { + _columnName = columnName; + _vector = vector; + _totalDocs = vector.getValueCount(); + _isSingleValue = !(vector instanceof ListVector); + _nextDocId = 0; + } + + @Override + public boolean hasNext() { + return _nextDocId < _totalDocs; + } + + @Override + @Nullable + public Object next() + throws IOException { + Object value = getValue(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public boolean isNextNull() + throws IOException { + return _vector.isNull(_nextDocId); + } + + @Override + public void skipNext() + throws IOException { + _nextDocId++; + } + + @Override + public boolean isSingleValue() { + return _isSingleValue; + } + + @Override + public boolean isInt() { + return _isSingleValue && (_vector instanceof IntVector || _vector instanceof BitVector); + } + + @Override + public boolean isLong() { + return _isSingleValue && _vector instanceof BigIntVector; + } + + @Override + public boolean isFloat() { + return _isSingleValue && _vector instanceof Float4Vector; + } + + @Override + public boolean isDouble() { + return _isSingleValue && _vector instanceof Float8Vector; + } + + @Override + public boolean isBigDecimal() { + return _isSingleValue && _vector instanceof DecimalVector; + } + + @Override + public boolean isString() { + return _isSingleValue && _vector instanceof VarCharVector; + } + + @Override + public boolean isBytes() { + return _isSingleValue && _vector instanceof VarBinaryVector; + } Review Comment: Fixed alongside the primitive predicates — `isString()` / `isBytes()` now check the list element vector type for multi-value columns (`VarCharVector` / `VarBinaryVector`). Same test, `testMultiValueColumnTypePredicates`. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowColumnReader.java: ########## @@ -0,0 +1,595 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.util.BitSet; +import javax.annotation.Nullable; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; + + +/** + * Column reader for Apache Arrow {@link FieldVector}. + * + * <p>Wraps a single Arrow {@link FieldVector} and exposes sequential and random-access + * read patterns conforming to the {@link ColumnReader} contract. The vector is owned by + * the enclosing {@link ArrowColumnReaderFactory} which is responsible for its lifecycle; + * closing this reader is a no-op for the underlying vector. + * + * <p>Supported Arrow types map to Pinot's stored types as follows: + * <ul> + * <li>{@code IntVector} / {@code BitVector} → {@code INT} (BitVector promoted to 0/1)</li> + * <li>{@code BigIntVector} → {@code LONG}</li> + * <li>{@code Float4Vector} → {@code FLOAT}</li> + * <li>{@code Float8Vector} → {@code DOUBLE}</li> + * <li>{@code DecimalVector} → {@code BIG_DECIMAL}</li> + * <li>{@code VarCharVector} → {@code STRING}</li> + * <li>{@code VarBinaryVector} → {@code BYTES}</li> + * <li>{@code ListVector} of the above → multi-value variant</li> + * </ul> + * + * <p>The list above applies to the typed primitive accessors only ({@link #getInt}, + * {@link #getString}, {@link #getIntMV}, ...). Complex types (Map, Struct, Union, ...) + * and temporal types are still readable via the generic {@link #getValue(int)} / + * {@link #next()} accessors, which delegate to {@link ArrowToPinotTypeConverter} and + * return the same canonical JDK types as the row-major path ({@link + * ArrowRecordExtractor}) — e.g. {@code Map<String, Object>} for Struct / Map, + * {@code Object[]} for List variants, {@code LocalDate} / {@code LocalTime} / + * {@code Timestamp} for temporal types. + * + * <p>This class is not thread-safe. + */ +public class ArrowColumnReader implements ColumnReader { + + private final String _columnName; + private final FieldVector _vector; + private final int _totalDocs; + private final boolean _isSingleValue; + + private int _nextDocId; + + /** + * Construct an ArrowColumnReader for the given vector. + * + * @param columnName Pinot column name + * @param vector Arrow field vector backing this column + */ + public ArrowColumnReader(String columnName, FieldVector vector) { + _columnName = columnName; + _vector = vector; + _totalDocs = vector.getValueCount(); + _isSingleValue = !(vector instanceof ListVector); + _nextDocId = 0; + } + + @Override + public boolean hasNext() { + return _nextDocId < _totalDocs; + } + + @Override + @Nullable + public Object next() + throws IOException { + Object value = getValue(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public boolean isNextNull() + throws IOException { + return _vector.isNull(_nextDocId); + } + + @Override + public void skipNext() + throws IOException { + _nextDocId++; + } + + @Override + public boolean isSingleValue() { + return _isSingleValue; + } + + @Override + public boolean isInt() { + return _isSingleValue && (_vector instanceof IntVector || _vector instanceof BitVector); + } + + @Override + public boolean isLong() { + return _isSingleValue && _vector instanceof BigIntVector; + } + + @Override + public boolean isFloat() { + return _isSingleValue && _vector instanceof Float4Vector; + } + + @Override + public boolean isDouble() { + return _isSingleValue && _vector instanceof Float8Vector; + } + + @Override + public boolean isBigDecimal() { + return _isSingleValue && _vector instanceof DecimalVector; + } + + @Override + public boolean isString() { + return _isSingleValue && _vector instanceof VarCharVector; + } + + @Override + public boolean isBytes() { + return _isSingleValue && _vector instanceof VarBinaryVector; + } + + @Override + public int nextInt() + throws IOException { + int value = getInt(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public long nextLong() + throws IOException { + long value = getLong(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public float nextFloat() + throws IOException { + float value = getFloat(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public double nextDouble() + throws IOException { + double value = getDouble(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public BigDecimal nextBigDecimal() + throws IOException { + BigDecimal value = getBigDecimal(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public String nextString() + throws IOException { + String value = getString(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public byte[] nextBytes() + throws IOException { + byte[] value = getBytes(_nextDocId); + _nextDocId++; + return value; + } + + @Override + public MultiValueResult<int[]> nextIntMV() + throws IOException { + MultiValueResult<int[]> result = getIntMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public MultiValueResult<long[]> nextLongMV() + throws IOException { + MultiValueResult<long[]> result = getLongMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public MultiValueResult<float[]> nextFloatMV() + throws IOException { + MultiValueResult<float[]> result = getFloatMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public MultiValueResult<double[]> nextDoubleMV() + throws IOException { + MultiValueResult<double[]> result = getDoubleMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public BigDecimal[] nextBigDecimalMV() + throws IOException { + BigDecimal[] result = getBigDecimalMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public String[] nextStringMV() + throws IOException { + String[] result = getStringMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public byte[][] nextBytesMV() + throws IOException { + byte[][] result = getBytesMV(_nextDocId); + _nextDocId++; + return result; + } + + @Override + public void rewind() + throws IOException { + _nextDocId = 0; + } + + @Override + public String getColumnName() { + return _columnName; + } + + @Override + public int getTotalDocs() { + return _totalDocs; + } + + @Override + public boolean isNull(int docId) { + checkBounds(docId); + return _vector.isNull(docId); + } + + @Override + public int getInt(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof IntVector) { + return ((IntVector) _vector).get(docId); + } + if (_vector instanceof BitVector) { + return ((BitVector) _vector).get(docId); + } + throw typeMismatch("INT"); + } + + @Override + public long getLong(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof BigIntVector) { + return ((BigIntVector) _vector).get(docId); + } + throw typeMismatch("LONG"); + } + + @Override + public float getFloat(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof Float4Vector) { + return ((Float4Vector) _vector).get(docId); + } + throw typeMismatch("FLOAT"); + } + + @Override + public double getDouble(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof Float8Vector) { + return ((Float8Vector) _vector).get(docId); + } + throw typeMismatch("DOUBLE"); + } + + @Override + public BigDecimal getBigDecimal(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof DecimalVector) { + return ((DecimalVector) _vector).getObject(docId); + } + throw typeMismatch("BIG_DECIMAL"); + } + + @Override + public String getString(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof VarCharVector) { + byte[] bytes = ((VarCharVector) _vector).get(docId); + return bytes == null ? null : new String(bytes, StandardCharsets.UTF_8); + } + throw typeMismatch("STRING"); + } + + @Override + public byte[] getBytes(int docId) + throws IOException { + checkBounds(docId); + if (_vector instanceof VarBinaryVector) { + return ((VarBinaryVector) _vector).get(docId); + } + throw typeMismatch("BYTES"); + } + + @Override + public Object getValue(int docId) + throws IOException { + checkBounds(docId); + Object value = _vector.getObject(docId); + if (value == null) { + return null; + } + // Delegate Arrow → Pinot type conversion to the shared utility extracted from + // ArrowRecordExtractor. Returns canonical JDK types: String for Utf8 / LargeUtf8 (unwrapped + // from Arrow's Text), Object[] for List variants (with recursive element conversion), + // LocalDate / LocalTime / Timestamp for temporal types, etc. + return ArrowToPinotTypeConverter.toPinotValue(_vector.getField(), value, false); + } Review Comment: Fixed — `extractRawTimeValues` is now threaded through the factory configs (`CONFIG_EXTRACT_RAW_TIME_VALUES`, mirroring `ArrowRecordExtractorConfig.EXTRACT_RAW_TIME_VALUES`) into the reader, and `getValue` honors it. New test `testExtractRawTimeValuesSurfacesRawEpoch` reads the same TIMESTAMP fixture with the flag on (raw epoch `long`) and off (canonical `Timestamp`). ########## pinot-plugins/pinot-input-format/pinot-arrow/src/test/java/org/apache/pinot/plugin/inputformat/arrow/ArrowFileColumnReaderFactoryTest.java: ########## @@ -0,0 +1,456 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.ipc.ArrowFileWriter; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Unit tests for {@link ArrowColumnReaderFactory} and {@link ArrowColumnReader}. + * + * <p>Each test materialises a small Arrow IPC file on disk with a known fixture and verifies + * that the factory can read it back column-by-column using all three documented patterns: + * generic sequential {@code next()}, typed sequential {@code nextInt() / nextLong() / ...}, + * and random-access {@code getInt(docId) / getLong(docId) / ...}. + */ +public class ArrowFileColumnReaderFactoryTest { + + private static final int ROWS = 8; + private Path _tempDir; + + @BeforeMethod + public void setUp() + throws IOException { + _tempDir = Files.createTempDirectory("arrow-column-reader-test"); + } + + @AfterMethod + public void tearDown() + throws IOException { + if (_tempDir != null) { + try (var stream = Files.walk(_tempDir)) { + stream.sorted((a, b) -> b.compareTo(a)).forEach(p -> { + try { + Files.deleteIfExists(p); + } catch (IOException ignored) { + // best effort + } + }); + } + } + } + + @Test + public void testSequentialReadAllPrimitiveTypes() + throws IOException { + File arrowFile = writePrimitiveFixture("primitive.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + + Map<String, ColumnReader> readers = factory.getAllColumnReaders(); + assertEquals(readers.size(), 6); + for (ColumnReader r : readers.values()) { + assertEquals(r.getTotalDocs(), ROWS); + assertTrue(r.isSingleValue()); + } + + ColumnReader intReader = readers.get("intCol"); + ColumnReader longReader = readers.get("longCol"); + ColumnReader floatReader = readers.get("floatCol"); + ColumnReader doubleReader = readers.get("doubleCol"); + ColumnReader stringReader = readers.get("stringCol"); + ColumnReader bytesReader = readers.get("bytesCol"); + + assertTrue(intReader.isInt()); + assertTrue(longReader.isLong()); + assertTrue(floatReader.isFloat()); + assertTrue(doubleReader.isDouble()); + assertTrue(stringReader.isString()); + assertTrue(bytesReader.isBytes()); + + // Pattern 2: typed sequential read with null handling + for (int i = 0; i < ROWS; i++) { + if (i == 3) { + assertTrue(intReader.isNextNull(), "Row 3 INT should be null"); + intReader.skipNext(); + assertTrue(longReader.isNextNull(), "Row 3 LONG should be null"); + longReader.skipNext(); + assertTrue(stringReader.isNextNull(), "Row 3 STRING should be null"); + stringReader.skipNext(); + floatReader.nextFloat(); + doubleReader.nextDouble(); + bytesReader.nextBytes(); + } else { + assertFalse(intReader.isNextNull()); + assertEquals(intReader.nextInt(), i * 10); + assertEquals(longReader.nextLong(), (long) i * 100); + assertEquals(floatReader.nextFloat(), i + 0.5f, 0.0001f); + assertEquals(doubleReader.nextDouble(), i + 0.25, 0.0001); + assertEquals(stringReader.nextString(), "s_" + i); + assertEquals(new String(bytesReader.nextBytes()), "b_" + i); + } Review Comment: Fixed — all byte/string conversions in the tests now pass `StandardCharsets.UTF_8` explicitly. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/test/java/org/apache/pinot/plugin/inputformat/arrow/ArrowFileColumnReaderFactoryTest.java: ########## @@ -0,0 +1,456 @@ +/** + * 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.plugin.inputformat.arrow; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.ipc.ArrowFileWriter; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.MultiValueResult; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Unit tests for {@link ArrowColumnReaderFactory} and {@link ArrowColumnReader}. + * + * <p>Each test materialises a small Arrow IPC file on disk with a known fixture and verifies + * that the factory can read it back column-by-column using all three documented patterns: + * generic sequential {@code next()}, typed sequential {@code nextInt() / nextLong() / ...}, + * and random-access {@code getInt(docId) / getLong(docId) / ...}. + */ +public class ArrowFileColumnReaderFactoryTest { + + private static final int ROWS = 8; + private Path _tempDir; + + @BeforeMethod + public void setUp() + throws IOException { + _tempDir = Files.createTempDirectory("arrow-column-reader-test"); + } + + @AfterMethod + public void tearDown() + throws IOException { + if (_tempDir != null) { + try (var stream = Files.walk(_tempDir)) { + stream.sorted((a, b) -> b.compareTo(a)).forEach(p -> { + try { + Files.deleteIfExists(p); + } catch (IOException ignored) { + // best effort + } + }); + } + } + } + + @Test + public void testSequentialReadAllPrimitiveTypes() + throws IOException { + File arrowFile = writePrimitiveFixture("primitive.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + + Map<String, ColumnReader> readers = factory.getAllColumnReaders(); + assertEquals(readers.size(), 6); + for (ColumnReader r : readers.values()) { + assertEquals(r.getTotalDocs(), ROWS); + assertTrue(r.isSingleValue()); + } + + ColumnReader intReader = readers.get("intCol"); + ColumnReader longReader = readers.get("longCol"); + ColumnReader floatReader = readers.get("floatCol"); + ColumnReader doubleReader = readers.get("doubleCol"); + ColumnReader stringReader = readers.get("stringCol"); + ColumnReader bytesReader = readers.get("bytesCol"); + + assertTrue(intReader.isInt()); + assertTrue(longReader.isLong()); + assertTrue(floatReader.isFloat()); + assertTrue(doubleReader.isDouble()); + assertTrue(stringReader.isString()); + assertTrue(bytesReader.isBytes()); + + // Pattern 2: typed sequential read with null handling + for (int i = 0; i < ROWS; i++) { + if (i == 3) { + assertTrue(intReader.isNextNull(), "Row 3 INT should be null"); + intReader.skipNext(); + assertTrue(longReader.isNextNull(), "Row 3 LONG should be null"); + longReader.skipNext(); + assertTrue(stringReader.isNextNull(), "Row 3 STRING should be null"); + stringReader.skipNext(); + floatReader.nextFloat(); + doubleReader.nextDouble(); + bytesReader.nextBytes(); + } else { + assertFalse(intReader.isNextNull()); + assertEquals(intReader.nextInt(), i * 10); + assertEquals(longReader.nextLong(), (long) i * 100); + assertEquals(floatReader.nextFloat(), i + 0.5f, 0.0001f); + assertEquals(doubleReader.nextDouble(), i + 0.25, 0.0001); + assertEquals(stringReader.nextString(), "s_" + i); + assertEquals(new String(bytesReader.nextBytes()), "b_" + i); + } + } + + // Verify rewind enables a second pass. + intReader.rewind(); + stringReader.rewind(); + assertTrue(intReader.hasNext()); + assertEquals(intReader.nextInt(), 0); + assertEquals(stringReader.nextString(), "s_0"); + } + } + + @Test + public void testRandomAccessByDocId() + throws IOException { + File arrowFile = writePrimitiveFixture("random.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + ColumnReader intReader = factory.getColumnReader("intCol"); + ColumnReader longReader = factory.getColumnReader("longCol"); + + // Pattern 3: read out of order. + assertEquals(intReader.getInt(5), 50); + assertEquals(intReader.getInt(0), 0); + assertEquals(intReader.getInt(7), 70); + assertTrue(intReader.isNull(3)); + assertEquals(longReader.getLong(2), 200L); + assertEquals(longReader.getValue(3), null); + + // Sequential read is unaffected by prior random-access reads on the same reader. + assertTrue(intReader.hasNext()); + assertEquals(intReader.nextInt(), 0); + } + } + + @Test + public void testGenericNextHandlesNulls() + throws IOException { + File arrowFile = writePrimitiveFixture("generic.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + ColumnReader stringReader = factory.getColumnReader("stringCol"); + + int nullCount = 0; + int nonNullCount = 0; + while (stringReader.hasNext()) { + Object value = stringReader.next(); + if (value == null) { + nullCount++; + } else { + nonNullCount++; + } + } + assertEquals(nullCount, 1); + assertEquals(nonNullCount, ROWS - 1); + } + } + + @Test + public void testMultiValueIntColumn() + throws IOException { + File arrowFile = writeMultiValueIntFixture("mv-int.arrow"); + org.apache.pinot.spi.data.Schema schema = new org.apache.pinot.spi.data.Schema.SchemaBuilder() + .setSchemaName("mv") + .addMultiValueDimension("intArr", FieldSpec.DataType.INT) + .build(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + ColumnReader reader = factory.getColumnReader("intArr"); + assertNotNull(reader); + assertFalse(reader.isSingleValue()); + assertEquals(reader.getTotalDocs(), 3); + + MultiValueResult<int[]> doc0 = reader.getIntMV(0); + assertEquals(doc0.getValues(), new int[]{1, 2, 3}); + assertFalse(doc0.hasNulls()); + + MultiValueResult<int[]> doc1 = reader.getIntMV(1); + assertEquals(doc1.getValues().length, 2); + assertTrue(doc1.hasNulls()); + assertTrue(doc1.isNull(1)); + assertFalse(doc1.isNull(0)); + assertEquals(doc1.getValues()[0], 4); + + MultiValueResult<int[]> doc2 = reader.getIntMV(2); + assertEquals(doc2.getValues().length, 0); + } + } + + @Test + public void testInitSubsetOfColumns() + throws IOException { + File arrowFile = writePrimitiveFixture("subset.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema, java.util.Set.of("intCol", "stringCol"), Collections.emptyMap()); + Map<String, ColumnReader> readers = factory.getAllColumnReaders(); + assertEquals(readers.size(), 2); + assertTrue(readers.containsKey("intCol")); + assertTrue(readers.containsKey("stringCol")); + assertNull(factory.getColumnReader("longCol")); + } + } + + @Test + public void testMultiBatchConcatenation() + throws IOException { + File arrowFile = writeMultiBatchFixture("multi-batch.arrow", 3, 4); + org.apache.pinot.spi.data.Schema schema = new org.apache.pinot.spi.data.Schema.SchemaBuilder() + .setSchemaName("multiBatch") + .addSingleValueDimension("seqCol", FieldSpec.DataType.INT) + .build(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema); + ColumnReader reader = factory.getColumnReader("seqCol"); + assertNotNull(reader); + assertEquals(reader.getTotalDocs(), 12, "Expected 3 batches of 4 rows = 12 docs total"); + + // Sequential traversal yields the global doc order. + for (int i = 0; i < 12; i++) { + assertEquals(reader.nextInt(), i); + } + assertFalse(reader.hasNext()); + + // Random access across batch boundaries. + reader.rewind(); + assertEquals(reader.getInt(0), 0); + assertEquals(reader.getInt(3), 3); // last row of batch 0 + assertEquals(reader.getInt(4), 4); // first row of batch 1 + assertEquals(reader.getInt(7), 7); // last row of batch 1 + assertEquals(reader.getInt(8), 8); // first row of batch 2 + assertEquals(reader.getInt(11), 11); // last row overall + } + } + + @Test + public void testGetAvailableColumnsIncludesAllSourceColumns() + throws IOException { + File arrowFile = writePrimitiveFixture("available.arrow"); + org.apache.pinot.spi.data.Schema schema = primitiveSchema(); + + try (ArrowFileColumnReaderFactory factory = new ArrowFileColumnReaderFactory(arrowFile)) { + factory.init(schema, java.util.Set.of("intCol"), Collections.emptyMap()); + // getAvailableColumns reflects the source schema, not the requested subset. + assertTrue(factory.getAvailableColumns().containsAll( + Arrays.asList("intCol", "longCol", "floatCol", "doubleCol", "stringCol", "bytesCol"))); + } + } + + // ===== Fixture builders ===== + + private org.apache.pinot.spi.data.Schema primitiveSchema() { + org.apache.pinot.spi.data.Schema schema = new org.apache.pinot.spi.data.Schema(); + schema.setSchemaName("primitive"); + schema.addField(new DimensionFieldSpec("intCol", FieldSpec.DataType.INT, true)); + schema.addField(new DimensionFieldSpec("longCol", FieldSpec.DataType.LONG, true)); + schema.addField(new DimensionFieldSpec("floatCol", FieldSpec.DataType.FLOAT, true)); + schema.addField(new DimensionFieldSpec("doubleCol", FieldSpec.DataType.DOUBLE, true)); + schema.addField(new DimensionFieldSpec("stringCol", FieldSpec.DataType.STRING, true)); + schema.addField(new DimensionFieldSpec("bytesCol", FieldSpec.DataType.BYTES, true)); + return schema; + } + + private File writePrimitiveFixture(String fileName) + throws IOException { + File out = _tempDir.resolve(fileName).toFile(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Field intField = new Field("intCol", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field longField = new Field("longCol", FieldType.nullable(new ArrowType.Int(64, true)), null); + Field floatField = + new Field("floatCol", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), null); + Field doubleField = + new Field("doubleCol", FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), null); + Field stringField = new Field("stringCol", FieldType.nullable(new ArrowType.Utf8()), null); + Field bytesField = new Field("bytesCol", FieldType.nullable(new ArrowType.Binary()), null); + Schema schema = + new Schema(Arrays.asList(intField, longField, floatField, doubleField, stringField, bytesField)); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + IntVector intVec = (IntVector) root.getVector("intCol"); + BigIntVector longVec = (BigIntVector) root.getVector("longCol"); + Float4Vector floatVec = (Float4Vector) root.getVector("floatCol"); + Float8Vector doubleVec = (Float8Vector) root.getVector("doubleCol"); + VarCharVector stringVec = (VarCharVector) root.getVector("stringCol"); + VarBinaryVector bytesVec = (VarBinaryVector) root.getVector("bytesCol"); + + intVec.allocateNew(ROWS); + longVec.allocateNew(ROWS); + floatVec.allocateNew(ROWS); + doubleVec.allocateNew(ROWS); + stringVec.allocateNew(ROWS * 8, ROWS); + bytesVec.allocateNew(ROWS * 8, ROWS); + + for (int i = 0; i < ROWS; i++) { + if (i == 3) { + intVec.setNull(i); + longVec.setNull(i); + stringVec.setNull(i); + } else { + intVec.set(i, i * 10); + longVec.set(i, (long) i * 100); + stringVec.set(i, ("s_" + i).getBytes()); + } + floatVec.set(i, i + 0.5f); + doubleVec.set(i, i + 0.25); + bytesVec.set(i, ("b_" + i).getBytes()); + } + intVec.setValueCount(ROWS); + longVec.setValueCount(ROWS); + floatVec.setValueCount(ROWS); + doubleVec.setValueCount(ROWS); + stringVec.setValueCount(ROWS); + bytesVec.setValueCount(ROWS); + root.setRowCount(ROWS); + + writeIpc(out, root); + } + } + return out; + } + + private File writeMultiBatchFixture(String fileName, int batchCount, int rowsPerBatch) + throws IOException { + File out = _tempDir.resolve(fileName).toFile(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Field seqField = new Field("seqCol", FieldType.nullable(new ArrowType.Int(32, true)), null); + Schema schema = new Schema(Collections.singletonList(seqField)); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + FileOutputStream fos = new FileOutputStream(out); + FileChannel channel = fos.getChannel(); + ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) { + writer.start(); + IntVector seqVec = (IntVector) root.getVector("seqCol"); + + int globalSeq = 0; + for (int b = 0; b < batchCount; b++) { + seqVec.allocateNew(rowsPerBatch); + for (int i = 0; i < rowsPerBatch; i++) { + seqVec.set(i, globalSeq++); + } + seqVec.setValueCount(rowsPerBatch); + root.setRowCount(rowsPerBatch); + writer.writeBatch(); + } + writer.end(); + } + } + return out; + } + + private File writeMultiValueIntFixture(String fileName) + throws IOException { + File out = _tempDir.resolve(fileName).toFile(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Field elementField = + new Field("element", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field listField = new Field("intArr", FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList(elementField)); + Schema schema = new Schema(Collections.singletonList(listField)); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + ListVector listVec = (ListVector) root.getVector("intArr"); + listVec.allocateNew(); + UnionListWriter writer = listVec.getWriter(); + + writer.startList(); + writer.setPosition(0); + writer.startList(); + writer.integer().writeInt(1); + writer.integer().writeInt(2); + writer.integer().writeInt(3); + writer.endList(); Review Comment: Cleaned up — the fixture now follows the documented `setPosition(i) → startList → writes → endList` pattern. The original leading `startList()` was redundant-but-benign on a fresh writer; the tests still pass. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
