| /* | |
| * 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.util.Preconditions; | |
| import org.apache.arrow.vector.BaseValueVector; | |
| import org.apache.arrow.vector.BitVectorHelper; | |
| import org.apache.arrow.vector.FieldVector; | |
| import org.apache.arrow.vector.ValueIterableVector; | |
| import org.apache.arrow.vector.ValueVector; | |
| import org.apache.arrow.vector.complex.AbstractStructVector; | |
| import org.apache.arrow.vector.complex.ListVector; | |
| 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; | |
| import org.apache.arrow.vector.util.TransferPair; | |
| import java.util.Arrays; | |
| import java.util.stream.Collectors; | |
| <@pp.dropOutputFile /> | |
| <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/DenseUnionVector.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.util.CommonUtil; | |
| import org.apache.arrow.memory.util.hash.ArrowBufHasher; | |
| import org.apache.arrow.memory.util.hash.SimpleHasher; | |
| 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.ipc.message.ArrowFieldNode; | |
| import org.apache.arrow.vector.BaseValueVector; | |
| import org.apache.arrow.vector.ValueIterableVector; | |
| import org.apache.arrow.vector.util.OversizedAllocationException; | |
| import org.apache.arrow.util.Preconditions; | |
| import static org.apache.arrow.vector.types.UnionMode.Dense; | |
| /* | |
| * 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, DenseUnionVector 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 DenseUnionVector extends AbstractContainerVector implements FieldVector, ValueIterableVector<Object> { | |
| int valueCount; | |
| NonNullableStructVector internalStruct; | |
| private ArrowBuf typeBuffer; | |
| private ArrowBuf offsetBuffer; | |
| /** | |
| * The key is type Id, and the value is vector. | |
| */ | |
| private ValueVector[] childVectors = new ValueVector[Byte.MAX_VALUE + 1]; | |
| /** | |
| * The index is the type id, and the value is the type field. | |
| */ | |
| private Field[] typeFields = new Field[Byte.MAX_VALUE + 1]; | |
| /** | |
| * The index is the index into the typeFields array, and the value is the logical field id. | |
| */ | |
| private byte[] typeMapFields = new byte[Byte.MAX_VALUE + 1]; | |
| /** | |
| * The next type id to allocate. | |
| */ | |
| private byte nextTypeId = 0; | |
| private FieldReader reader; | |
| private long typeBufferAllocationSizeInBytes; | |
| private long offsetBufferAllocationSizeInBytes; | |
| private final FieldType fieldType; | |
| public static final byte TYPE_WIDTH = 1; | |
| public static final byte OFFSET_WIDTH = 4; | |
| private static final FieldType INTERNAL_STRUCT_TYPE = new FieldType(/*nullable*/ false, | |
| ArrowType.Struct.INSTANCE, /*dictionary*/ null, /*metadata*/ null); | |
| public static DenseUnionVector empty(String name, BufferAllocator allocator) { | |
| FieldType fieldType = FieldType.notNullable(new ArrowType.Union( | |
| UnionMode.Dense, null)); | |
| return new DenseUnionVector(name, allocator, fieldType, null); | |
| } | |
| public DenseUnionVector(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; | |
| this.offsetBuffer = allocator.getEmpty(); | |
| this.offsetBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; | |
| } | |
| public BufferAllocator getAllocator() { | |
| return allocator; | |
| } | |
| @Override | |
| public MinorType getMinorType() { | |
| return MinorType.DENSEUNION; | |
| } | |
| @Override | |
| public void initializeChildrenFromFields(List<Field> children) { | |
| for (Field field : children) { | |
| byte typeId = registerNewTypeId(field); | |
| FieldVector vector = (FieldVector) internalStruct.add(field.getName(), field.getFieldType()); | |
| vector.initializeChildrenFromFields(field.getChildren()); | |
| childVectors[typeId] = vector; | |
| } | |
| } | |
| @Override | |
| public List<FieldVector> getChildrenFromFields() { | |
| return internalStruct.getChildrenFromFields(); | |
| } | |
| @Override | |
| public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> ownBuffers) { | |
| if (ownBuffers.size() != 2) { | |
| throw new IllegalArgumentException("Illegal buffer count for dense union with type " + getField().getFieldType() + | |
| ", expected " + 2 + ", got: " + ownBuffers.size()); | |
| } | |
| ArrowBuf buffer = ownBuffers.get(0); | |
| typeBuffer.getReferenceManager().release(); | |
| typeBuffer = buffer.getReferenceManager().retain(buffer, allocator); | |
| typeBufferAllocationSizeInBytes = typeBuffer.capacity(); | |
| buffer = ownBuffers.get(1); | |
| offsetBuffer.getReferenceManager().release(); | |
| offsetBuffer = buffer.getReferenceManager().retain(buffer, allocator); | |
| offsetBufferAllocationSizeInBytes = offsetBuffer.capacity(); | |
| this.valueCount = fieldNode.getLength(); | |
| } | |
| @Override | |
| public List<ArrowBuf> getFieldBuffers() { | |
| List<ArrowBuf> result = new ArrayList<>(2); | |
| setReaderAndWriterIndex(); | |
| result.add(typeBuffer); | |
| result.add(offsetBuffer); | |
| return result; | |
| } | |
| private void setReaderAndWriterIndex() { | |
| typeBuffer.readerIndex(0); | |
| typeBuffer.writerIndex(valueCount * TYPE_WIDTH); | |
| offsetBuffer.readerIndex(0); | |
| offsetBuffer.writerIndex((long) valueCount * OFFSET_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 | |
| */ | |
| @Override | |
| @Deprecated | |
| public List<BufferBacked> getFieldInnerVectors() { | |
| throw new UnsupportedOperationException("There are no inner vectors. Use geFieldBuffers"); | |
| } | |
| private String fieldName(byte typeId, MinorType type) { | |
| return type.name().toLowerCase() + typeId; | |
| } | |
| private FieldType fieldType(MinorType type) { | |
| return FieldType.nullable(type.getType()); | |
| } | |
| public synchronized byte registerNewTypeId(Field field) { | |
| if (nextTypeId == typeFields.length) { | |
| throw new IllegalStateException("Dense union vector support at most " + | |
| typeFields.length + " relative types. Please use union of union instead"); | |
| } | |
| byte typeId = nextTypeId; | |
| if (this.fieldType != null) { | |
| int[] typeIds = ((ArrowType.Union) this.fieldType.getType()).getTypeIds(); | |
| if (typeIds != null) { | |
| int thisTypeId = typeIds[nextTypeId]; | |
| if (thisTypeId > Byte.MAX_VALUE) { | |
| throw new IllegalStateException("Dense union vector types must be bytes. " + thisTypeId + " is too large"); | |
| } | |
| typeId = (byte) thisTypeId; | |
| } | |
| } | |
| typeFields[typeId] = field; | |
| typeMapFields[nextTypeId] = typeId; | |
| this.nextTypeId += 1; | |
| return typeId; | |
| } | |
| private <T extends FieldVector> T addOrGet(byte typeId, MinorType minorType, Class<T> c) { | |
| return internalStruct.addOrGet(fieldName(typeId, minorType), fieldType(minorType), c); | |
| } | |
| private <T extends FieldVector> T addOrGet(byte typeId, MinorType minorType, ArrowType arrowType, Class<T> c) { | |
| return internalStruct.addOrGet(fieldName(typeId, minorType), FieldType.nullable(arrowType), c); | |
| } | |
| @Override | |
| public long getOffsetBufferAddress() { | |
| return offsetBuffer.memoryAddress(); | |
| } | |
| @Override | |
| public long getDataBufferAddress() { | |
| throw new UnsupportedOperationException(); | |
| } | |
| @Override | |
| public long getValidityBufferAddress() { | |
| throw new UnsupportedOperationException(); | |
| } | |
| @Override | |
| public ArrowBuf getValidityBuffer() { throw new UnsupportedOperationException(); } | |
| @Override | |
| public ArrowBuf getOffsetBuffer() { return offsetBuffer; } | |
| public ArrowBuf getTypeBuffer() { return typeBuffer; } | |
| @Override | |
| public ArrowBuf getDataBuffer() { throw new UnsupportedOperationException(); } | |
| public StructVector getStruct(byte typeId) { | |
| StructVector structVector = typeId < 0 ? null : (StructVector) childVectors[typeId]; | |
| if (structVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| structVector = addOrGet(typeId, MinorType.STRUCT, StructVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| structVector.allocateNew(); | |
| childVectors[typeId] = structVector; | |
| 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")> | |
| public ${name}Vector get${name}Vector(byte typeId<#if minor.class?starts_with("Decimal")>, ArrowType arrowType</#if>) { | |
| ValueVector vector = typeId < 0 ? null : childVectors[typeId]; | |
| if (vector == null) { | |
| int vectorCount = internalStruct.size(); | |
| vector = addOrGet(typeId, MinorType.${name?upper_case}<#if minor.class?starts_with("Decimal")>, arrowType</#if>, ${name}Vector.class); | |
| childVectors[typeId] = vector; | |
| if (internalStruct.size() > vectorCount) { | |
| vector.allocateNew(); | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return (${name}Vector) vector; | |
| } | |
| </#if> | |
| </#list> | |
| </#list> | |
| public ListVector getList(byte typeId) { | |
| ListVector listVector = typeId < 0 ? null : (ListVector) childVectors[typeId]; | |
| if (listVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| listVector = addOrGet(typeId, MinorType.LIST, ListVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| listVector.allocateNew(); | |
| childVectors[typeId] = listVector; | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return listVector; | |
| } | |
| public MapVector getMap(byte typeId) { | |
| MapVector mapVector = typeId < 0 ? null : (MapVector) childVectors[typeId]; | |
| if (mapVector == null) { | |
| int vectorCount = internalStruct.size(); | |
| mapVector = addOrGet(typeId, MinorType.MAP, MapVector.class); | |
| if (internalStruct.size() > vectorCount) { | |
| mapVector.allocateNew(); | |
| childVectors[typeId] = mapVector; | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| } | |
| } | |
| return mapVector; | |
| } | |
| public byte getTypeId(int index) { | |
| return typeBuffer.getByte(index * TYPE_WIDTH); | |
| } | |
| public ValueVector getVectorByType(byte typeId) { | |
| return typeId < 0 ? null : childVectors[typeId]; | |
| } | |
| @Override | |
| public void allocateNew() throws OutOfMemoryException { | |
| /* new allocation -- clear the current buffers */ | |
| clear(); | |
| internalStruct.allocateNew(); | |
| try { | |
| allocateTypeBuffer(); | |
| allocateOffsetBuffer(); | |
| } 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(); | |
| allocateOffsetBuffer(); | |
| } catch (Exception e) { | |
| clear(); | |
| return false; | |
| } | |
| return true; | |
| } | |
| private void allocateTypeBuffer() { | |
| typeBuffer = allocator.buffer(typeBufferAllocationSizeInBytes); | |
| typeBuffer.readerIndex(0); | |
| setNegative(0, typeBuffer.capacity()); | |
| } | |
| private void allocateOffsetBuffer() { | |
| offsetBuffer = allocator.buffer(offsetBufferAllocationSizeInBytes); | |
| offsetBuffer.readerIndex(0); | |
| offsetBuffer.setZero(0, offsetBuffer.capacity()); | |
| } | |
| @Override | |
| public void reAlloc() { | |
| internalStruct.reAlloc(); | |
| reallocTypeBuffer(); | |
| reallocOffsetBuffer(); | |
| } | |
| public int getOffset(int index) { | |
| return offsetBuffer.getInt((long) index * OFFSET_WIDTH); | |
| } | |
| 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((int)newAllocationSize); | |
| newBuf.setBytes(0, typeBuffer, 0, currentBufferCapacity); | |
| typeBuffer.getReferenceManager().release(1); | |
| typeBuffer = newBuf; | |
| typeBufferAllocationSizeInBytes = (int)newAllocationSize; | |
| setNegative(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); | |
| } | |
| private void reallocOffsetBuffer() { | |
| final long currentBufferCapacity = offsetBuffer.capacity(); | |
| long newAllocationSize = currentBufferCapacity * 2; | |
| if (newAllocationSize == 0) { | |
| if (offsetBufferAllocationSizeInBytes > 0) { | |
| newAllocationSize = offsetBufferAllocationSizeInBytes; | |
| } else { | |
| newAllocationSize = BaseValueVector.INITIAL_VALUE_ALLOCATION * OFFSET_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((int) newAllocationSize); | |
| newBuf.setBytes(0, offsetBuffer, 0, currentBufferCapacity); | |
| newBuf.setZero(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); | |
| offsetBuffer.getReferenceManager().release(1); | |
| offsetBuffer = newBuf; | |
| offsetBufferAllocationSizeInBytes = (int) newAllocationSize; | |
| } | |
| @Override | |
| public void setInitialCapacity(int numRecords) { } | |
| @Override | |
| public int getValueCapacity() { | |
| long capacity = getTypeBufferValueCapacity(); | |
| long offsetCapacity = getOffsetBufferValueCapacity(); | |
| if (offsetCapacity < capacity) { | |
| capacity = offsetCapacity; | |
| } | |
| long structCapacity = internalStruct.getValueCapacity(); | |
| if (structCapacity < capacity) { | |
| structCapacity = capacity; | |
| } | |
| return (int) capacity; | |
| } | |
| @Override | |
| public void close() { | |
| clear(); | |
| } | |
| @Override | |
| public void clear() { | |
| valueCount = 0; | |
| typeBuffer.getReferenceManager().release(); | |
| typeBuffer = allocator.getEmpty(); | |
| offsetBuffer.getReferenceManager().release(); | |
| offsetBuffer = allocator.getEmpty(); | |
| internalStruct.clear(); | |
| } | |
| @Override | |
| public void reset() { | |
| valueCount = 0; | |
| setNegative(0, typeBuffer.capacity()); | |
| offsetBuffer.setZero(0, offsetBuffer.capacity()); | |
| internalStruct.reset(); | |
| } | |
| @Override | |
| public Field getField() { | |
| int childCount = (int) Arrays.stream(typeFields).filter(field -> field != null).count(); | |
| List<org.apache.arrow.vector.types.pojo.Field> childFields = new ArrayList<>(childCount); | |
| int[] typeIds = new int[childCount]; | |
| for (int i = 0; i < typeFields.length; i++) { | |
| if (typeFields[i] != null) { | |
| int curIdx = childFields.size(); | |
| typeIds[curIdx] = i; | |
| childFields.add(typeFields[i]); | |
| } | |
| } | |
| FieldType fieldType; | |
| if (this.fieldType == null) { | |
| fieldType = FieldType.nullable(new ArrowType.Union(Dense, typeIds)); | |
| } else { | |
| final UnionMode mode = UnionMode.Dense; | |
| 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.DenseUnionVector.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.DenseUnionVector.TransferImpl(field, allocator, callBack); | |
| } | |
| @Override | |
| public TransferPair makeTransferPair(ValueVector target) { | |
| return new TransferImpl((DenseUnionVector) target); | |
| } | |
| @Override | |
| public void copyFrom(int inIndex, int outIndex, ValueVector from) { | |
| Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); | |
| DenseUnionVector fromCast = (DenseUnionVector) from; | |
| int inOffset = fromCast.offsetBuffer.getInt((long) inIndex * OFFSET_WIDTH); | |
| fromCast.getReader().setPosition(inOffset); | |
| int outOffset = offsetBuffer.getInt((long) outIndex * OFFSET_WIDTH); | |
| getWriter().setPosition(outOffset); | |
| ComplexCopier.copy(fromCast.reader, writer); | |
| } | |
| @Override | |
| public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { | |
| copyFrom(inIndex, outIndex, from); | |
| } | |
| public FieldVector addVector(byte typeId, FieldVector v) { | |
| final String name = v.getName().isEmpty() ? fieldName(typeId, 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); | |
| childVectors[typeId] = newVector; | |
| if (callBack != null) { | |
| callBack.doWork(); | |
| } | |
| return newVector; | |
| } | |
| private class TransferImpl implements TransferPair { | |
| private final TransferPair[] internalTransferPairs = new TransferPair[nextTypeId]; | |
| private final DenseUnionVector to; | |
| public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) { | |
| to = new DenseUnionVector(name, allocator, null, callBack); | |
| internalStruct.makeTransferPair(to.internalStruct); | |
| createTransferPairs(); | |
| } | |
| public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) { | |
| to = new DenseUnionVector(field.getName(), allocator, null, callBack); | |
| internalStruct.makeTransferPair(to.internalStruct); | |
| createTransferPairs(); | |
| } | |
| public TransferImpl(DenseUnionVector to) { | |
| this.to = to; | |
| internalStruct.makeTransferPair(to.internalStruct); | |
| createTransferPairs(); | |
| } | |
| private void createTransferPairs() { | |
| for (int i = 0; i < nextTypeId; i++) { | |
| ValueVector srcVec = internalStruct.getVectorById(i); | |
| ValueVector dstVec = to.internalStruct.getVectorById(i); | |
| to.typeFields[i] = typeFields[i]; | |
| to.typeMapFields[i] = typeMapFields[i]; | |
| to.childVectors[i] = dstVec; | |
| internalTransferPairs[i] = srcVec.makeTransferPair(dstVec); | |
| } | |
| } | |
| @Override | |
| public void transfer() { | |
| to.clear(); | |
| ReferenceManager refManager = typeBuffer.getReferenceManager(); | |
| to.typeBuffer = refManager.transferOwnership(typeBuffer, to.allocator).getTransferredBuffer(); | |
| refManager = offsetBuffer.getReferenceManager(); | |
| to.offsetBuffer = refManager.transferOwnership(offsetBuffer, to.allocator).getTransferredBuffer(); | |
| for (int i = 0; i < nextTypeId; i++) { | |
| if (internalTransferPairs[i] != null) { | |
| internalTransferPairs[i].transfer(); | |
| to.childVectors[i] = internalTransferPairs[i].getTo(); | |
| } | |
| } | |
| to.valueCount = valueCount; | |
| clear(); | |
| } | |
| @Override | |
| public void splitAndTransfer(int startIndex, int length) { | |
| to.clear(); | |
| // transfer type buffer | |
| int startPoint = startIndex * TYPE_WIDTH; | |
| int sliceLength = length * TYPE_WIDTH; | |
| ArrowBuf slicedBuffer = typeBuffer.slice(startPoint, sliceLength); | |
| ReferenceManager refManager = slicedBuffer.getReferenceManager(); | |
| to.typeBuffer = refManager.transferOwnership(slicedBuffer, to.allocator).getTransferredBuffer(); | |
| // transfer offset buffer | |
| while (to.offsetBuffer.capacity() < (long) length * OFFSET_WIDTH) { | |
| to.reallocOffsetBuffer(); | |
| } | |
| int [] typeCounts = new int[nextTypeId]; | |
| int [] typeStarts = new int[nextTypeId]; | |
| for (int i = 0; i < typeCounts.length; i++) { | |
| typeCounts[i] = 0; | |
| typeStarts[i] = -1; | |
| } | |
| for (int i = startIndex; i < startIndex + length; i++) { | |
| byte typeId = typeBuffer.getByte(i); | |
| if (typeId >= 0) { | |
| to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, typeCounts[typeId]); | |
| typeCounts[typeId] += 1; | |
| if (typeStarts[typeId] == -1) { | |
| typeStarts[typeId] = offsetBuffer.getInt((long) i * OFFSET_WIDTH); | |
| } | |
| } | |
| } | |
| // transfer vector values | |
| for (int i = 0; i < nextTypeId; i++) { | |
| if (typeCounts[i] > 0 && typeStarts[i] != -1) { | |
| internalTransferPairs[i].splitAndTransfer(typeStarts[i], typeCounts[i]); | |
| to.childVectors[i] = internalTransferPairs[i].getTo(); | |
| } | |
| } | |
| to.setValueCount(length); | |
| } | |
| @Override | |
| public ValueVector getTo() { | |
| return to; | |
| } | |
| @Override | |
| public void copyValueSafe(int from, int to) { | |
| this.to.copyFrom(from, to, DenseUnionVector.this); | |
| } | |
| } | |
| @Override | |
| public FieldReader getReader() { | |
| if (reader == null) { | |
| reader = new DenseUnionReader(this); | |
| } | |
| return reader; | |
| } | |
| public FieldWriter getWriter() { | |
| if (writer == null) { | |
| writer = new DenseUnionWriter(this); | |
| } | |
| return writer; | |
| } | |
| @Override | |
| public int getBufferSize() { | |
| return this.getBufferSizeFor(this.valueCount); | |
| } | |
| @Override | |
| public int getBufferSizeFor(final int count) { | |
| if (count == 0) { | |
| return 0; | |
| } | |
| int[] counts = new int[Byte.MAX_VALUE + 1]; | |
| for (int i = 0; i < count; i++) { | |
| byte typeId = getTypeId(i); | |
| if (typeId != -1) { | |
| counts[typeId] += 1; | |
| } | |
| } | |
| long childBytes = 0; | |
| for (int typeId = 0; typeId < childVectors.length; typeId++) { | |
| ValueVector childVector = childVectors[typeId]; | |
| if (childVector != null) { | |
| childBytes += childVector.getBufferSizeFor(counts[typeId]); | |
| } | |
| } | |
| return (int) (count * TYPE_WIDTH + (long) count * OFFSET_WIDTH + childBytes); | |
| } | |
| @Override | |
| public ArrowBuf[] getBuffers(boolean clear) { | |
| List<ArrowBuf> list = new java.util.ArrayList<>(); | |
| setReaderAndWriterIndex(); | |
| if (getBufferSize() != 0) { | |
| list.add(typeBuffer); | |
| list.add(offsetBuffer); | |
| list.addAll(java.util.Arrays.asList(internalStruct.getBuffers(clear))); | |
| } | |
| if (clear) { | |
| valueCount = 0; | |
| typeBuffer.getReferenceManager().retain(); | |
| typeBuffer.close(); | |
| typeBuffer = allocator.getEmpty(); | |
| offsetBuffer.getReferenceManager().retain(); | |
| offsetBuffer.close(); | |
| offsetBuffer = allocator.getEmpty(); | |
| } | |
| return list.toArray(new ArrowBuf[list.size()]); | |
| } | |
| @Override | |
| public Iterator<ValueVector> iterator() { | |
| return internalStruct.iterator(); | |
| } | |
| private ValueVector getVector(int index) { | |
| byte typeId = typeBuffer.getByte(index * TYPE_WIDTH); | |
| return getVectorByType(typeId); | |
| } | |
| public Object getObject(int index) { | |
| ValueVector vector = getVector(index); | |
| if (vector != null) { | |
| int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); | |
| return vector.isNull(offset) ? null : vector.getObject(offset); | |
| } | |
| return null; | |
| } | |
| public void get(int index, DenseUnionHolder holder) { | |
| FieldReader reader = new DenseUnionReader(DenseUnionVector.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; | |
| } | |
| DenseUnionWriter writer; | |
| public void setValueCount(int valueCount) { | |
| this.valueCount = valueCount; | |
| while (valueCount > getTypeBufferValueCapacity()) { | |
| reallocTypeBuffer(); | |
| reallocOffsetBuffer(); | |
| } | |
| setChildVectorValueCounts(); | |
| } | |
| private void setChildVectorValueCounts() { | |
| int [] counts = new int[Byte.MAX_VALUE + 1]; | |
| for (int i = 0; i < this.valueCount; i++) { | |
| byte typeId = getTypeId(i); | |
| if (typeId != -1) { | |
| counts[typeId] += 1; | |
| } | |
| } | |
| for (int i = 0; i < nextTypeId; i++) { | |
| childVectors[typeMapFields[i]].setValueCount(counts[typeMapFields[i]]); | |
| } | |
| } | |
| public void setSafe(int index, DenseUnionHolder holder) { | |
| FieldReader reader = holder.reader; | |
| if (writer == null) { | |
| writer = new DenseUnionWriter(DenseUnionVector.this); | |
| } | |
| int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); | |
| MinorType type = reader.getMinorType(); | |
| writer.setPosition(offset); | |
| byte typeId = holder.typeId; | |
| 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")> | |
| case ${name?upper_case}: | |
| Nullable${name}Holder ${uncappedName}Holder = new Nullable${name}Holder(); | |
| reader.read(${uncappedName}Holder); | |
| setSafe(index, ${uncappedName}Holder); | |
| break; | |
| </#if> | |
| </#list> | |
| </#list> | |
| case STRUCT: | |
| case LIST: { | |
| setTypeId(index, typeId); | |
| 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")> | |
| public void setSafe(int index, Nullable${name}Holder holder) { | |
| while (index >= getOffsetBufferValueCapacity()) { | |
| reallocOffsetBuffer(); | |
| } | |
| byte typeId = getTypeId(index); | |
| ${name}Vector vector = get${name}Vector(typeId<#if minor.class?starts_with("Decimal")>, new ArrowType.Decimal(holder.precision, holder.scale, holder.WIDTH * 8)</#if>); | |
| int offset = vector.getValueCount(); | |
| vector.setValueCount(offset + 1); | |
| vector.setSafe(offset, holder); | |
| offsetBuffer.setInt((long) index * OFFSET_WIDTH, offset); | |
| } | |
| </#if> | |
| </#list> | |
| </#list> | |
| public void setTypeId(int index, byte typeId) { | |
| while (index >= getTypeBufferValueCapacity()) { | |
| reallocTypeBuffer(); | |
| } | |
| typeBuffer.setByte(index * TYPE_WIDTH , typeId); | |
| } | |
| private int getTypeBufferValueCapacity() { | |
| return (int) typeBuffer.capacity() / TYPE_WIDTH; | |
| } | |
| public void setOffset(int index, int offset) { | |
| while (index >= getOffsetBufferValueCapacity()) { | |
| reallocOffsetBuffer(); | |
| } | |
| offsetBuffer.setInt((long) index * OFFSET_WIDTH, offset); | |
| } | |
| private long getOffsetBufferValueCapacity() { | |
| return offsetBuffer.capacity() / OFFSET_WIDTH; | |
| } | |
| @Override | |
| public int hashCode(int index, ArrowBufHasher hasher) { | |
| if (isNull(index)) { | |
| return 0; | |
| } | |
| int offset = offsetBuffer.getInt((long) index * OFFSET_WIDTH); | |
| return getVector(index).hashCode(offset, hasher); | |
| } | |
| @Override | |
| public int hashCode(int index) { | |
| return hashCode(index, SimpleHasher.INSTANCE); | |
| } | |
| @Override | |
| public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) { | |
| return visitor.visit(this, value); | |
| } | |
| @Override | |
| public String getName() { | |
| return name; | |
| } | |
| private void setNegative(long start, long end) { | |
| for (long i = start;i < end; i++) { | |
| typeBuffer.setByte(i, -1); | |
| } | |
| } | |
| @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 DenseUnionVector, 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 DenseUnionVector."); | |
| } | |
| } |