pvary commented on code in PR #16871: URL: https://github.com/apache/iceberg/pull/16871#discussion_r3812879390
########## 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: Use new switch -- 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]
