Copilot commented on code in PR #18638: URL: https://github.com/apache/pinot/pull/18638#discussion_r3384164557
########## 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: If `populate()` throws (e.g., empty Arrow source triggers the `!anyBatchSeen` exception, or `loadNextBatch()` fails), the already-allocated accumulator vectors are leaked because they are never closed. This can pin significant off-heap memory in the Arrow allocator on failure paths. ########## 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: Javadoc references `InMemoryArrowColumnReaderFactory`, but no such class exists in this module; the helper is shared by `ArrowColumnReaderFactory` and `ArrowFileColumnReaderFactory`. Broken Javadoc links can fail doclint builds and are confusing for readers. ########## 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: Sequential iteration methods should throw `IllegalStateException` when `hasNext()` is false (matching `PinotSegmentColumnReaderImpl` / `DefaultValueColumnReader`). Current implementation can throw `IndexOutOfBoundsException` from Arrow vectors and `skipNext()` can silently advance past the end, which makes iterator misuse harder to diagnose and diverges from the established SPI behavior. ########## 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: `isInt()`/`isLong()`/`isFloat()`/`isDouble()`/`isBigDecimal()` currently return `false` for multi-value columns because they are gated on `_isSingleValue`. The SPI uses these methods for both SV and MV (see `DataTypeColumnTransformer`), so MV primitive/list columns will fall back to boxed `next()` paths even though `get*MV()`/`next*MV()` are implemented. These predicates should also return `true` when the reader is a `ListVector` whose *element* vector matches the expected type. ########## 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: Test uses platform-default charset when decoding bytes (`new String(byte[])`), which can make the test non-deterministic across environments. Specify UTF-8 explicitly (the fixture encodes ASCII/UTF-8 bytes). ########## 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: Same as the primitive predicates: `isString()` and `isBytes()` should return `true` for multi-value list columns when the underlying element vector is `VarCharVector` / `VarBinaryVector`. Currently they always return `false` for MV, forcing boxed paths unexpectedly. ########## 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: The multi-value list fixture calls `writer.startList()` twice before writing the first row, and the first call is not paired with an `endList()`. This is likely incorrect usage of `UnionListWriter` and can produce malformed vectors or hide bugs (depending on Arrow version). ########## 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: `getValue(int)` hard-codes `extractRawTimeValues=false`, so the column-major path cannot honor the `extractRawTimeValues` behavior supported by the row-major `ArrowRecordReaderConfig`/`ArrowRecordExtractorConfig`. If callers rely on raw temporal values (e.g. for schema/unit-preserving ingestion), row-major and column-major ingestion will diverge. ########## 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}. + * Review Comment: Class Javadoc says these are tests for `ArrowColumnReaderFactory`, but this file specifically tests the file-backed `ArrowFileColumnReaderFactory`. Updating the link avoids confusion when navigating tests. ########## pinot-plugins/pinot-input-format/pinot-arrow/src/main/java/org/apache/pinot/plugin/inputformat/arrow/ArrowFileColumnReaderFactory.java: ########## @@ -0,0 +1,214 @@ +/** + * 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.FileInputStream; +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ipc.ArrowFileReader; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.ColumnReader; +import org.apache.pinot.spi.data.readers.ColumnReaderFactory; + + +/** + * {@link ColumnReaderFactory} backed by an Apache Arrow IPC file on disk. + * + * <p>File-specialised convenience over {@link ArrowColumnReaderFactory}: this class opens a private + * {@link RootAllocator} sized by {@link #CONFIG_ALLOCATOR_LIMIT}, opens the file via + * {@link ArrowFileReader}, concatenates all record batches into per-column accumulators via the + * shared {@link ArrowAccumulators} helper, and closes the file, reader, and allocator on + * {@link #close}. Callers that already manage an {@link org.apache.arrow.vector.ipc.ArrowReader} + * and {@link org.apache.arrow.memory.BufferAllocator} themselves should use + * {@link ArrowColumnReaderFactory} directly. + * + * <p>Columns in the target schema that are absent from the Arrow file are reported via + * {@link #getAvailableColumns()}, and {@link #getColumnReader(String)} returns {@code null} for + * them; schema-evolution defaults are the columnar build driver's responsibility. Review Comment: Javadoc is backwards: `getAvailableColumns()` reports columns present in the Arrow file (source schema), not target-schema columns that are absent. As written it suggests the opposite and can confuse callers about schema evolution behavior. -- 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]
