thswlsqls commented on code in PR #16871:
URL: https://github.com/apache/iceberg/pull/16871#discussion_r3974046149
##########
arrow/src/main/java/org/apache/iceberg/arrow/vectorized/ArrowBatchReader.java:
##########
@@ -43,16 +50,41 @@ public final ColumnarBatch read(ColumnarBatch reuse, int
numRowsToRead) {
ColumnVector[] columnVectors = new ColumnVector[readers.length];
for (int i = 0; i < readers.length; i += 1) {
- vectorHolders[i] = readers[i].read(vectorHolders[i], numRowsToRead);
- int numRowsInVector = vectorHolders[i].numValues();
+ VectorHolder holder = readers[i].read(vectorHolders[i], numRowsToRead);
+ if (holder.isDummy()) {
+ // The column is not in the data file, so the reader returns the value
it should be read as
+ // instead of a vector. Build a vector holding that value because
callers read a batch as an
+ // Arrow VectorSchemaRoot, which has no way to represent a column
without a vector.
+ closeVector(vectorHolders[i]);
+ holder = constantHolder(fields.get(i), holder, numRowsToRead);
+ }
+
+ vectorHolders[i] = holder;
+ int numRowsInVector = holder.numValues();
Preconditions.checkState(
numRowsInVector == numRowsToRead,
"Number of rows in the vector %s didn't match expected %s ",
numRowsInVector,
numRowsToRead);
- // Handle null vector for constant case
- columnVectors[i] = new ColumnVector(vectorHolders[i]);
+ columnVectors[i] = new ColumnVector(holder);
}
return new ColumnarBatch(numRowsToRead, columnVectors);
}
+
+ private VectorHolder constantHolder(
+ Types.NestedField field, VectorHolder dummyHolder, int numRows) {
+ Object constant =
+ dummyHolder instanceof ConstantVectorHolder
+ ? ((ConstantVectorHolder<?>) dummyHolder).getConstant()
+ : null;
Review Comment:
@pvary Thanks — the constant path does carry it, but the null path doesn't,
so I kept `fields` here.
`ConstantVectorReader` builds its holder with
`constantHolder(icebergField(), …)`, while `VectorizedArrowReader.nulls()`
returns a shared `NullVectorReader` singleton whose `dummyHolder` leaves the
field null. That path covers optional columns with no default, and the
`rowIds`/`lastUpdated` fallbacks.
I can give the null reader its own field so the holder always has one, and
then drop `fields` from this class — want me to?
##########
arrow/src/main/java/org/apache/iceberg/arrow/vectorized/ArrowBatchReader.java:
##########
@@ -43,16 +50,41 @@ public final ColumnarBatch read(ColumnarBatch reuse, int
numRowsToRead) {
ColumnVector[] columnVectors = new ColumnVector[readers.length];
for (int i = 0; i < readers.length; i += 1) {
- vectorHolders[i] = readers[i].read(vectorHolders[i], numRowsToRead);
- int numRowsInVector = vectorHolders[i].numValues();
+ VectorHolder holder = readers[i].read(vectorHolders[i], numRowsToRead);
+ if (holder.isDummy()) {
+ // The column is not in the data file, so the reader returns the value
it should be read as
+ // instead of a vector. Build a vector holding that value because
callers read a batch as an
+ // Arrow VectorSchemaRoot, which has no way to represent a column
without a vector.
+ closeVector(vectorHolders[i]);
Review Comment:
@pvary Thanks — you're right, I'll reuse the vector instead of closing and
rebuilding it every batch.
Right now the constant vector is freed and reallocated per batch even when
`reuse` is non-null, so passing it buys nothing for constant columns. I'll keep
the one built for the column and refill it when the row count changes.
Let me know if you meant something else by it.
##########
arrow/src/main/java/org/apache/iceberg/arrow/vectorized/ConstantVectors.java:
##########
@@ -0,0 +1,215 @@
+/*
+ * 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.iceberg.arrow.vectorized;
+
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.UUID;
+import java.util.function.IntConsumer;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
+import org.apache.arrow.vector.DecimalVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.TimeMicroVector;
+import org.apache.arrow.vector.TimeStampMicroTZVector;
+import org.apache.arrow.vector.TimeStampMicroVector;
+import org.apache.arrow.vector.TimeStampNanoTZVector;
+import org.apache.arrow.vector.TimeStampNanoVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.iceberg.arrow.ArrowSchemaUtil;
+import org.apache.iceberg.types.Type;
+import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.ByteBuffers;
+import org.apache.iceberg.util.UUIDUtil;
+
+/**
+ * Builds an Arrow vector whose rows all hold the same value.
+ *
+ * <p>A column that is missing from a data file, either because it was added
by a schema change or
+ * because it has a default value, is read as a constant instead of from the
file. Spark represents
+ * such a column with a virtual vector that never allocates memory, but the
Arrow reader hands out
+ * {@link org.apache.arrow.vector.VectorSchemaRoot}, so it needs a real vector
holding the value
+ * repeated once per row.
+ */
+class ConstantVectors {
+
+ private ConstantVectors() {}
+
+ /**
+ * Creates a holder for a vector of {@code numRows} rows that all hold
{@code constant}.
+ *
+ * @param field the Iceberg field the vector is read for, used to derive the
Arrow type
+ * @param constant the value every row holds, or null to produce an all-null
vector
+ * @param numRows the number of rows in the vector
+ * @param allocator the allocator that owns the vector's memory
+ * @return a holder whose vector the caller is responsible for closing
+ */
+ static VectorHolder holder(
+ Types.NestedField field, Object constant, int numRows, BufferAllocator
allocator) {
+ FieldVector vector = createVector(field, constant, numRows, allocator);
+
+ NullabilityHolder nullabilityHolder = new NullabilityHolder(numRows);
+ if (constant == null) {
+ nullabilityHolder.setNulls(0, numRows);
+ } else {
+ nullabilityHolder.setNotNulls(0, numRows);
+ }
+
+ return VectorHolder.vectorHolder(vector, field, nullabilityHolder);
+ }
+
+ private static FieldVector createVector(
+ Types.NestedField field, Object constant, int numRows, BufferAllocator
allocator) {
+ FieldVector vector =
ArrowSchemaUtil.convert(field).createVector(allocator);
+ try {
+ vector.setInitialCapacity(numRows);
+ vector.allocateNew();
+
+ if (constant != null) {
+ IntConsumer setter = setter(vector, field.type(), constant);
+ for (int row = 0; row < numRows; row += 1) {
+ setter.accept(row);
+ }
+ }
+
+ // rows that were never set keep the validity bit allocateNew cleared,
so they read as null
+ vector.setValueCount(numRows);
+ return vector;
+ } catch (RuntimeException e) {
+ vector.close();
+ throw e;
+ }
+ }
+
+ private static IntConsumer setter(FieldVector vector, Type type, Object
constant) {
+ switch (type.typeId()) {
Review Comment:
@pvary Thanks — I'll convert it to the arrow form (`case … ->`).
--
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]
