| /* | |
| * 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. | |
| */ | |
| import org.apache.arrow.memory.ArrowBuf; | |
| import org.apache.arrow.memory.BufferAllocator; | |
| import org.apache.arrow.memory.ReferenceManager; | |
| import org.apache.arrow.memory.util.CommonUtil; | |
| import org.apache.arrow.memory.util.hash.ArrowBufHasher; | |
| import org.apache.arrow.util.Preconditions; | |
| import org.apache.arrow.vector.BaseValueVector; | |
| import org.apache.arrow.vector.BitVectorHelper; | |
| import org.apache.arrow.vector.DecimalVector; | |
| import org.apache.arrow.vector.FieldVector; | |
| import org.apache.arrow.vector.ValueVector; | |
| import org.apache.arrow.vector.complex.AbstractStructVector; | |
| import org.apache.arrow.vector.complex.NonNullableStructVector; | |
| import org.apache.arrow.vector.complex.StructVector; | |
| import org.apache.arrow.vector.compare.VectorVisitor; | |
| import org.apache.arrow.vector.types.Types; | |
| import org.apache.arrow.vector.types.UnionMode; | |
| import org.apache.arrow.vector.compare.RangeEqualsVisitor; | |
| 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.util.CallBack; | |
| import org.apache.arrow.vector.util.DataSizeRoundingUtil; | |
| <@pp.dropOutputFile /> | |
| <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/UnionVector.java" /> | |
| <#include "/@includes/license.ftl" /> | |
| package org.apache.arrow.vector.complex; | |
| <#include "/@includes/vv_imports.ftl" /> | |
| import java.util.ArrayList; | |
| import java.util.Collections; | |
| import java.util.Iterator; | |
| import org.apache.arrow.memory.ArrowBuf; | |
| import org.apache.arrow.memory.util.CommonUtil; | |
| import org.apache.arrow.vector.compare.VectorVisitor; | |
| import org.apache.arrow.vector.complex.impl.ComplexCopier; | |
| import org.apache.arrow.vector.util.CallBack; | |
| import org.apache.arrow.vector.util.ValueVectorUtility; | |
| import org.apache.arrow.vector.ipc.message.ArrowFieldNode; | |
| import org.apache.arrow.memory.util.ArrowBufPointer; | |
| import org.apache.arrow.memory.util.hash.ArrowBufHasher; | |
| import org.apache.arrow.vector.BaseValueVector; | |
| import org.apache.arrow.vector.util.OversizedAllocationException; | |
| import org.apache.arrow.util.Preconditions; | |
| import static org.apache.arrow.vector.types.UnionMode.Sparse; | |
| import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; | |
| import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt; | |
| <#function is_timestamp_tz type> | |
| <#return type?starts_with("TimeStamp") && type?ends_with("TZ")> | |
| </#function> | |
| /* | |
| * This class is generated using freemarker and the ${.template_name} template. | |
| */ | |
| @SuppressWarnings("unused") | |
| /** | |
| * A vector which can hold values of different types. It does so by using a StructVector which contains a vector for each | |
| * primitive type that is stored. StructVector is used in order to take advantage of its serialization/deserialization methods, | |
| * as well as the addOrGet method. | |
| * | |
| * For performance reasons, UnionVector stores a cached reference to each subtype vector, to avoid having to do the struct lookup | |
| * each time the vector is accessed. | |
| * Source code generated using FreeMarker template ${.template_name} | |
| */ | |
| public class UnionVector extends AbstractContainerVector implements FieldVector { | |
| int valueCount; | |
| NonNullableStructVector internalStruct; | |
| protected ArrowBuf typeBuffer; | |
| private StructVector structVector; | |
| private ListVector listVector; | |
| private ListViewVector listViewVector; | |
| private MapVector mapVector; | |
| private FieldReader reader; | |
| private int singleType = 0; | |
| private ValueVector singleVector; | |
| private int typeBufferAllocationSizeInBytes; | |
| private final FieldType fieldType; | |
| private final Field[] typeIds = new Field[Byte.MAX_VALUE + 1]; | |
| public static final byte TYPE_WIDTH = 1; | |
| private static final FieldType INTERNAL_STRUCT_TYPE = new FieldType(false /*nullable*/, | |
| ArrowType.Struct.INSTANCE, null /*dictionary*/, null /*metadata*/); | |
| public static UnionVector empty(String name, BufferAllocator allocator) { | |
| FieldType fieldType = FieldType.nullable(new ArrowType.Union( | |
| UnionMode.Sparse, null)); | |
| return new UnionVector(name, allocator, fieldType, null); | |
| } | |
| public UnionVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) { | |
| super(name, allocator, callBack); | |
| this.fieldType = fieldType; | |
| this.internalStruct = new NonNullableStructVector( | |
| "internal", | |
| allocator, | |
| INTERNAL_STRUCT_TYPE, | |
| callBack, | |
| AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE, | |
| false); | |
| this.typeBuffer = allocator.getEmpty(); | |
| this.typeBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH; | |
| } | |
| public BufferAllocator getAllocator() { | |
| return allocator; | |
| } | |
| @Override | |
| public MinorType getMinorType() { | |
| return MinorType.UNION; | |
| } | |
| @Override | |
| public void initializeChildrenFromFields(List<Field> children) { | |
| int count = 0; | |
| for (Field child: children) { | |
| int typeId = Types.getMinorTypeForArrowType(child.getType()).ordinal(); | |
| if (this.fieldType != null) { | |
| int[] typeIds = ((ArrowType.Union)this.fieldType.getType()).getTypeIds(); | |
| if (typeIds != null) { | |
| typeId = typeIds[count++]; | |
| } | |
| } | |
| typeIds[typeId] = child; | |
| } | |
| internalStruct.initializeChildrenFromFields(children); | |
| } | |
| @Override | |
| public List<FieldVector> getChildrenFromFields() { | |
| return internalStruct.getChildrenFromFields(); | |
| } | |
| @Override | |
| public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> ownBuffers) { | |
| if (ownBuffers.size() != 1) { | |
| throw new IllegalArgumentException("Illegal buffer count, expected 1, got: " + ownBuffers.size()); | |
| } | |
| ArrowBuf buffer = ownBuffers.get(0); | |
| typeBuffer.getReferenceManager().release(); | |
| typeBuffer = buffer.getReferenceManager().retain(buffer, allocator); | |
| typeBufferAllocationSizeInBytes = checkedCastToInt(typeBuffer.capacity()); | |
| this.valueCount = fieldNode.getLength(); | |
| } | |
| @Override | |
| public List<ArrowBuf> getFieldBuffers() { | |
| List<ArrowBuf> result = new ArrayList<>(1); | |
| setReaderAndWriterIndex(); | |
| result.add(typeBuffer); | |
| return result; | |
| } | |
| private void setReaderAndWriterIndex() { | |
| typeBuffer.readerIndex(0); | |
| typeBuffer.writerIndex(valueCount * TYPE_WIDTH); | |
| } | |
| /** | |
| * Get the inner vectors. | |
| * | |
| * @deprecated This API will be removed as the current implementations no longer support inner vectors. | |
| * | |
| * @return the inner vectors for this field as defined by the TypeLayout | |
| */ | |
| @Deprecated | |
| @Override | |
| public List<BufferBacked> getFieldInnerVectors() { | |
| throw new UnsupportedOperationException("There are no inner vectors. Use geFieldBuffers"); | |
| } | |
| private String fieldName(MinorType type) { | |
| return type.name().toLowerCase(); | |
| } | |
| private FieldType fieldType(MinorType type) { | |
| return FieldType.nullable(type.getType()); | |
| } | |
| private <T extends FieldVector> T addOrGet(Types.MinorType minorType, Class<T> c) { | |
| return addOrGet(null, minorType, c); | |
| } | |
| private <T extends FieldVector> T addOrGet(String name, Types.MinorType minorType, ArrowType arrowType, Class<T> c) { | |
| return internalStruct.addOrGet(name == null ? fieldName(minorType) : name, FieldType.nullable(arrowType), c); | |
| } | |
| private <T extends FieldVector> T addOrGet(String name, Types.MinorType minorType, Class<T> c) { | |
| return internalStruct.addOrGet(name == null ? fieldName(minorType) : name, fieldType(minorType), c); | |
| } | |
| @Override | |
| public long getValidityBufferAddress() { | |
| throw new UnsupportedOperationException(); | |
| } | |
| public long getTypeBufferAddress() { | |
| return typeBuffer.memoryAddress(); | |
| } | |
| @Override | |
| public long getDataBufferAddress() { | |
| throw new UnsupportedOperationException(); | |
| } | |
| @Override | |
| public long getOffsetBufferAddress() { | |
| throw new UnsupportedOperationException(); | |
| } | |
| public ArrowBuf getTypeBuffer() { | |
| return typeBuffer; | |
| } | |
| @Override | |
| public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } | |
| @Override | |
| public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } | |
| @Override | |
| public ArrowBuf getOffsetBuffer() { throw new UnsupportedOperationException(); } | |
| public StructVector getStruct() { | |
| if (structVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| structVector = addOrGet(MinorType.STRUCT, StructVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| structVector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return structVector; | |
| } | |
| <#list vv.types as type> | |
| <#list type.minor as minor> | |
| <#assign name = minor.class?cap_first /> | |
| <#assign fields = minor.fields!type.fields /> | |
| <#assign uncappedName = name?uncap_first/> | |
| <#assign lowerCaseName = name?lower_case/> | |
| <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| private ${name}Vector ${uncappedName}Vector; | |
| <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| public ${name}Vector get${name}Vector() { | |
| if (${uncappedName}Vector == null) { | |
| ${uncappedName}Vector = internalStruct.getChild(fieldName(MinorType.${name?upper_case}), ${name}Vector.class); | |
| if (${uncappedName}Vector == null) { | |
| throw new IllegalArgumentException("No ${name} present. Provide ArrowType argument to create a new vector"); | |
| } | |
| } | |
| return ${uncappedName}Vector; | |
| } | |
| public ${name}Vector get${name}Vector(ArrowType arrowType) { | |
| return get${name}Vector(null, arrowType); | |
| } | |
| public ${name}Vector get${name}Vector(String name, ArrowType arrowType) { | |
| if (${uncappedName}Vector == null) { | |
| int vectorCount = internalStruct.size(); | |
| ${uncappedName}Vector = addOrGet(name, MinorType.${name?upper_case}, arrowType, ${name}Vector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| ${uncappedName}Vector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return ${uncappedName}Vector; | |
| } | |
| <#else> | |
| public ${name}Vector get${name}Vector() { | |
| return get${name}Vector(null); | |
| } | |
| public ${name}Vector get${name}Vector(String name) { | |
| if (${uncappedName}Vector == null) { | |
| int vectorCount = internalStruct.size(); | |
| ${uncappedName}Vector = addOrGet(name, MinorType.${name?upper_case}, ${name}Vector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| ${uncappedName}Vector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return ${uncappedName}Vector; | |
| } | |
| </#if> | |
| </#if> | |
| </#list> | |
| </#list> | |
| public ListVector getList() { | |
| if (listVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| listVector = addOrGet(MinorType.LIST, ListVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| listVector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return listVector; | |
| } | |
| public ListViewVector getListView() { | |
| if (listViewVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| listViewVector = addOrGet(MinorType.LISTVIEW, ListViewVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| listViewVector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return listViewVector; | |
| } | |
| public MapVector getMap() { | |
| if (mapVector == null) { | |
| throw new IllegalArgumentException("No map present. Provide ArrowType argument to create a new vector"); | |
| } | |
| return mapVector; | |
| } | |
| public MapVector getMap(ArrowType arrowType) { | |
| return getMap(null, arrowType); | |
| } | |
| public MapVector getMap(String name, ArrowType arrowType) { | |
| if (mapVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| mapVector = addOrGet(name, MinorType.MAP, arrowType, MapVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| mapVector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return mapVector; | |
| } | |
| private ExtensionTypeVector extensionVector; | |
| public ExtensionTypeVector getExtension(ArrowType arrowType) { | |
| if (extensionVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| extensionVector = addOrGet(null, MinorType.EXTENSIONTYPE, arrowType, ExtensionTypeVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| extensionVector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return extensionVector; | |
| } | |
| public int getTypeValue(int index) { | |
| return typeBuffer.getByte(index * TYPE_WIDTH); | |
| } | |
| @Override | |
| public void allocateNew() throws OutOfMemoryException { | |
| /* new allocation -- clear the current buffers */ | |
| clear(); | |
| internalStruct.allocateNew(); | |
| try { | |
| allocateTypeBuffer(); | |
| } catch (Exception e) { | |
| clear(); | |
| throw e; | |
| } | |
| } | |
| @Override | |
| public boolean allocateNewSafe() { | |
| /* new allocation -- clear the current buffers */ | |
| clear(); | |
| boolean safe = internalStruct.allocateNewSafe(); | |
| if (!safe) { return false; } | |
| try { | |
| allocateTypeBuffer(); | |
| } catch (Exception e) { | |
| clear(); | |
| return false; | |
| } | |
| return true; | |
| } | |
| private void allocateTypeBuffer() { | |
| typeBuffer = allocator.buffer(typeBufferAllocationSizeInBytes); | |
| typeBuffer.readerIndex(0); | |
| typeBuffer.setZero(0, typeBuffer.capacity()); | |
| } | |
| @Override | |
| public void reAlloc() { | |
| internalStruct.reAlloc(); | |
| reallocTypeBuffer(); | |
| } | |
| private void reallocTypeBuffer() { | |
| final long currentBufferCapacity = typeBuffer.capacity(); | |
| long newAllocationSize = currentBufferCapacity * 2; | |
| if (newAllocationSize == 0) { | |
| if (typeBufferAllocationSizeInBytes > 0) { | |
| newAllocationSize = typeBufferAllocationSizeInBytes; | |
| } else { | |
| newAllocationSize = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH * 2; | |
| } | |
| } | |
| newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); | |
| assert newAllocationSize >= 1; | |
| if (newAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { | |
| throw new OversizedAllocationException("Unable to expand the buffer"); | |
| } | |
| final ArrowBuf newBuf = allocator.buffer(checkedCastToInt(newAllocationSize)); | |
| newBuf.setBytes(0, typeBuffer, 0, currentBufferCapacity); | |
| newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); | |
| typeBuffer.getReferenceManager().release(1); | |
| typeBuffer = newBuf; | |
| typeBufferAllocationSizeInBytes = (int)newAllocationSize; | |
| } | |
| @Override | |
| public void setInitialCapacity(int numRecords) { } | |
| @Override | |
| public int getValueCapacity() { | |
| return Math.min(getTypeBufferValueCapacity(), internalStruct.getValueCapacity()); | |
| } | |
| @Override | |
| public void close() { | |
| clear(); | |
| } | |
| @Override | |
| public void clear() { | |
| valueCount = 0; | |
| typeBuffer.getReferenceManager().release(); | |
| typeBuffer = allocator.getEmpty(); | |
| internalStruct.clear(); | |
| } | |
| @Override | |
| public void reset() { | |
| valueCount = 0; | |
| typeBuffer.setZero(0, typeBuffer.capacity()); | |
| internalStruct.reset(); | |
| } | |
| @Override | |
| public Field getField() { | |
| List<org.apache.arrow.vector.types.pojo.Field> childFields = new ArrayList<>(); | |
| List<FieldVector> children = internalStruct.getChildren(); | |
| int[] typeIds = new int[children.size()]; | |
| for (ValueVector v : children) { | |
| typeIds[childFields.size()] = v.getMinorType().ordinal(); | |
| childFields.add(v.getField()); | |
| } | |
| FieldType fieldType; | |
| if (this.fieldType == null) { | |
| fieldType = FieldType.nullable(new ArrowType.Union(Sparse, typeIds)); | |
| } else { | |
| final UnionMode mode = ((ArrowType.Union)this.fieldType.getType()).getMode(); | |
| fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds), | |
| this.fieldType.getDictionary(), this.fieldType.getMetadata()); | |
| } | |
| return new Field(name, fieldType, childFields); | |
| } | |
| @Override | |
| public TransferPair getTransferPair(BufferAllocator allocator) { | |
| return getTransferPair(name, allocator); | |
| } | |
| @Override | |
| public TransferPair getTransferPair(String ref, BufferAllocator allocator) { | |
| return getTransferPair(ref, allocator, null); | |
| } | |
| @Override | |
| public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { | |
| return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(ref, allocator, callBack); | |
| } | |
| @Override | |
| public TransferPair getTransferPair(Field field, BufferAllocator allocator) { | |
| return getTransferPair(field, allocator, null); | |
| } | |
| @Override | |
| public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { | |
| return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(field, allocator, callBack); | |
| } | |
| @Override | |
| public TransferPair makeTransferPair(ValueVector target) { | |
| return new TransferImpl((UnionVector) target); | |
| } | |
| @Override | |
| public void copyFrom(int inIndex, int outIndex, ValueVector from) { | |
| Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); | |
| UnionVector fromCast = (UnionVector) from; | |
| fromCast.getReader().setPosition(inIndex); | |
| getWriter().setPosition(outIndex); | |
| ComplexCopier.copy(fromCast.reader, writer); | |
| } | |
| @Override | |
| public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { | |
| copyFrom(inIndex, outIndex, from); | |
| } | |
| public FieldVector addVector(FieldVector v) { | |
| final String name = v.getName().isEmpty() ? fieldName(v.getMinorType()) : v.getName(); | |
| Preconditions.checkState(internalStruct.getChild(name) == null, String.format("%s vector already exists", name)); | |
| final FieldVector newVector = internalStruct.addOrGet(name, v.getField().getFieldType(), v.getClass()); | |
| v.makeTransferPair(newVector).transfer(); | |
| internalStruct.putChild(name, newVector); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| return newVector; | |
| } | |
| /** | |
| * Directly put a vector to internalStruct without creating a new one with same type. | |
| */ | |
| public void directAddVector(FieldVector v) { | |
| String name = fieldName(v.getMinorType()); | |
| Preconditions.checkState(internalStruct.getChild(name) == null, String.format("%s vector already exists", name)); | |
| internalStruct.putChild(name, v); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| private class TransferImpl implements TransferPair { | |
| private final TransferPair internalStructVectorTransferPair; | |
| private final UnionVector to; | |
| public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { | |
| to = new UnionVector(name, allocator, /* field type */ null, callBack); | |
| internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); | |
| } | |
| public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { | |
| to = new UnionVector(field.getName(), allocator, null, callBack); | |
| internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); | |
| } | |
| public TransferImpl(UnionVector to) { | |
| this.to = to; | |
| internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct); | |
| } | |
| @Override | |
| public void transfer() { | |
| to.clear(); | |
| ReferenceManager refManager = typeBuffer.getReferenceManager(); | |
| to.typeBuffer = refManager.transferOwnership(typeBuffer, to.allocator).getTransferredBuffer(); | |
| internalStructVectorTransferPair.transfer(); | |
| to.valueCount = valueCount; | |
| clear(); | |
| } | |
| @Override | |
| public void splitAndTransfer(int startIndex, int length) { | |
| Preconditions.checkArgument(startIndex >= 0 && length >= 0 && startIndex + length <= valueCount, | |
| "Invalid parameters startIndex: %s, length: %s for valueCount: %s", startIndex, length, valueCount); | |
| to.clear(); | |
| internalStructVectorTransferPair.splitAndTransfer(startIndex, length); | |
| final int startPoint = startIndex * TYPE_WIDTH; | |
| final int sliceLength = length * TYPE_WIDTH; | |
| final ArrowBuf slicedBuffer = typeBuffer.slice(startPoint, sliceLength); | |
| final ReferenceManager refManager = slicedBuffer.getReferenceManager(); | |
| to.typeBuffer = refManager.transferOwnership(slicedBuffer, to.allocator).getTransferredBuffer(); | |
| to.setValueCount(length); | |
| } | |
| @Override | |
| public ValueVector getTo() { | |
| return to; | |
| } | |
| @Override | |
| public void copyValueSafe(int from, int to) { | |
| this.to.copyFrom(from, to, UnionVector.this); | |
| } | |
| } | |
| @Override | |
| public FieldReader getReader() { | |
| if (reader == null) { | |
| reader = new UnionReader(this); | |
| } | |
| return reader; | |
| } | |
| public FieldWriter getWriter() { | |
| if (writer == null) { | |
| writer = new UnionWriter(this); | |
| } | |
| return writer; | |
| } | |
| @Override | |
| public int getBufferSize() { | |
| if (valueCount == 0) { return 0; } | |
| return (valueCount * TYPE_WIDTH) + internalStruct.getBufferSize(); | |
| } | |
| @Override | |
| public int getBufferSizeFor(final int valueCount) { | |
| if (valueCount == 0) { | |
| return 0; | |
| } | |
| long bufferSize = 0; | |
| for (final ValueVector v : (Iterable<ValueVector>) this) { | |
| bufferSize += v.getBufferSizeFor(valueCount); | |
| } | |
| return (int) bufferSize + (valueCount * TYPE_WIDTH); | |
| } | |
| @Override | |
| public ArrowBuf[] getBuffers(boolean clear) { | |
| List<ArrowBuf> list = new java.util.ArrayList<>(); | |
| setReaderAndWriterIndex(); | |
| if (getBufferSize() != 0) { | |
| list.add(typeBuffer); | |
| list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); | |
| } | |
| if (clear) { | |
| valueCount = 0; | |
| typeBuffer.getReferenceManager().retain(); | |
| typeBuffer.getReferenceManager().release(); | |
| typeBuffer = allocator.getEmpty(); | |
| } | |
| return list.toArray(new ArrowBuf[list.size()]); | |
| } | |
| @Override | |
| public Iterator<ValueVector> iterator() { | |
| return internalStruct.iterator(); | |
| } | |
| public ValueVector getVector(int index) { | |
| return getVector(index, null); | |
| } | |
| public ValueVector getVector(int index, ArrowType arrowType) { | |
| int type = typeBuffer.getByte(index * TYPE_WIDTH); | |
| return getVectorByType(type, arrowType); | |
| } | |
| public ValueVector getVectorByType(int typeId) { | |
| return getVectorByType(typeId, null); | |
| } | |
| public ValueVector getVectorByType(int typeId, ArrowType arrowType) { | |
| Field type = typeIds[typeId]; | |
| Types.MinorType minorType; | |
| String name = null; | |
| if (type == null) { | |
| minorType = Types.MinorType.values()[typeId]; | |
| } else { | |
| minorType = Types.getMinorTypeForArrowType(type.getType()); | |
| name = type.getName(); | |
| } | |
| switch (minorType) { | |
| case NULL: | |
| return null; | |
| <#list vv.types as type> | |
| <#list type.minor as minor> | |
| <#assign name = minor.class?cap_first /> | |
| <#assign fields = minor.fields!type.fields /> | |
| <#assign uncappedName = name?uncap_first/> | |
| <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| case ${name?upper_case}: | |
| return get${name}Vector(name<#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary">, arrowType</#if>); | |
| </#if> | |
| </#list> | |
| </#list> | |
| case STRUCT: | |
| return getStruct(); | |
| case LIST: | |
| return getList(); | |
| case LISTVIEW: | |
| return getListView(); | |
| case MAP: | |
| return getMap(name, arrowType); | |
| case EXTENSIONTYPE: | |
| return getExtension(arrowType); | |
| default: | |
| throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]); | |
| } | |
| } | |
| public Object getObject(int index) { | |
| ValueVector vector = getVector(index); | |
| if (vector != null) { | |
| return vector.isNull(index) ? null : vector.getObject(index); | |
| } | |
| return null; | |
| } | |
| public byte[] get(int index) { | |
| return null; | |
| } | |
| public void get(int index, ComplexHolder holder) { | |
| } | |
| public void get(int index, UnionHolder holder) { | |
| FieldReader reader = new UnionReader(UnionVector.this); | |
| reader.setPosition(index); | |
| holder.reader = reader; | |
| } | |
| public int getValueCount() { | |
| return valueCount; | |
| } | |
| /** | |
| * IMPORTANT: Union types always return non null as there is no validity buffer. | |
| * | |
| * To check validity correctly you must check the underlying vector. | |
| */ | |
| public boolean isNull(int index) { | |
| return false; | |
| } | |
| @Override | |
| public int getNullCount() { | |
| return 0; | |
| } | |
| public int isSet(int index) { | |
| return isNull(index) ? 0 : 1; | |
| } | |
| UnionWriter writer; | |
| public void setValueCount(int valueCount) { | |
| this.valueCount = valueCount; | |
| while (valueCount > getTypeBufferValueCapacity()) { | |
| reallocTypeBuffer(); | |
| } | |
| internalStruct.setValueCount(valueCount); | |
| } | |
| public void setSafe(int index, UnionHolder holder) { | |
| setSafe(index, holder, null); | |
| } | |
| public void setSafe(int index, UnionHolder holder, ArrowType arrowType) { | |
| FieldReader reader = holder.reader; | |
| if (writer == null) { | |
| writer = new UnionWriter(UnionVector.this); | |
| } | |
| writer.setPosition(index); | |
| MinorType type = reader.getMinorType(); | |
| switch (type) { | |
| <#list vv.types as type> | |
| <#list type.minor as minor> | |
| <#assign name = minor.class?cap_first /> | |
| <#assign fields = minor.fields!type.fields /> | |
| <#assign uncappedName = name?uncap_first/> | |
| <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| case ${name?upper_case}: | |
| Nullable${name}Holder ${uncappedName}Holder = new Nullable${name}Holder(); | |
| reader.read(${uncappedName}Holder); | |
| <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| setSafe(index, ${uncappedName}Holder, arrowType); | |
| <#else> | |
| setSafe(index, ${uncappedName}Holder); | |
| </#if> | |
| break; | |
| </#if> | |
| </#list> | |
| </#list> | |
| case STRUCT: { | |
| ComplexCopier.copy(reader, writer); | |
| break; | |
| } | |
| case LIST: { | |
| ComplexCopier.copy(reader, writer); | |
| break; | |
| } | |
| default: | |
| throw new UnsupportedOperationException(); | |
| } | |
| } | |
| <#list vv.types as type> | |
| <#list type.minor as minor> | |
| <#assign name = minor.class?cap_first /> | |
| <#assign fields = minor.fields!type.fields /> | |
| <#assign uncappedName = name?uncap_first/> | |
| <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> | |
| public void setSafe(int index, Nullable${name}Holder holder, ArrowType arrowType) { | |
| setType(index, MinorType.${name?upper_case}); | |
| get${name}Vector(null, arrowType).setSafe(index, holder); | |
| } | |
| <#else> | |
| public void setSafe(int index, Nullable${name}Holder holder) { | |
| setType(index, MinorType.${name?upper_case}); | |
| get${name}Vector(null).setSafe(index, holder); | |
| } | |
| </#if> | |
| </#if> | |
| </#list> | |
| </#list> | |
| public void setType(int index, MinorType type) { | |
| while (index >= getTypeBufferValueCapacity()) { | |
| reallocTypeBuffer(); | |
| } | |
| typeBuffer.setByte(index * TYPE_WIDTH , (byte) type.ordinal()); | |
| } | |
| private int getTypeBufferValueCapacity() { | |
| return capAtMaxInt(typeBuffer.capacity() / TYPE_WIDTH); | |
| } | |
| @Override | |
| public int hashCode(int index) { | |
| return hashCode(index, null); | |
| } | |
| @Override | |
| public int hashCode(int index, ArrowBufHasher hasher) { | |
| ValueVector vec = getVector(index); | |
| if (vec == null) { | |
| return ArrowBufPointer.NULL_HASH_CODE; | |
| } | |
| return vec.hashCode(index, hasher); | |
| } | |
| @Override | |
| public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) { | |
| return visitor.visit(this, value); | |
| } | |
| @Override | |
| public String getName() { | |
| return name; | |
| } | |
| @Override | |
| public String toString() { | |
| return ValueVectorUtility.getToString(this, 0, getValueCount()); | |
| } | |
| @Override | |
| public <T extends FieldVector> T addOrGet(String name, FieldType fieldType, Class<T> clazz) { | |
| return internalStruct.addOrGet(name, fieldType, clazz); | |
| } | |
| @Override | |
| public <T extends FieldVector> T getChild(String name, Class<T> clazz) { | |
| return internalStruct.getChild(name, clazz); | |
| } | |
| @Override | |
| public VectorWithOrdinal getChildVectorWithOrdinal(String name) { | |
| return internalStruct.getChildVectorWithOrdinal(name); | |
| } | |
| @Override | |
| public int size() { | |
| return internalStruct.size(); | |
| } | |
| @Override | |
| public void setInitialCapacity(int valueCount, double density) { | |
| for (final ValueVector vector : internalStruct) { | |
| if (vector instanceof DensityAwareVector) { | |
| ((DensityAwareVector) vector).setInitialCapacity(valueCount, density); | |
| } else { | |
| vector.setInitialCapacity(valueCount); | |
| } | |
| } | |
| } | |
| /** | |
| * Set the element at the given index to null. For UnionVector, it throws an UnsupportedOperationException | |
| * as nulls are not supported at the top level and isNull() always returns false. | |
| * | |
| * @param index position of element | |
| * @throws UnsupportedOperationException whenever invoked | |
| */ | |
| @Override | |
| public void setNull(int index) { | |
| throw new UnsupportedOperationException("The method setNull() is not supported on UnionVector."); | |
| } | |
| } |