diff --git a/vector/src/main/codegen/templates/DenseUnionVector.java b/vector/src/main/codegen/templates/DenseUnionVector.java index e9a9a2878b..8dbf34ea87 100644 --- a/vector/src/main/codegen/templates/DenseUnionVector.java +++ b/vector/src/main/codegen/templates/DenseUnionVector.java @@ -74,7 +74,7 @@ /* * This class is generated using freemarker and the ${.template_name} template. */ -@SuppressWarnings("unused") +@SuppressWarnings({"unused", "static"}) /** @@ -102,6 +102,12 @@ public class DenseUnionVector extends AbstractContainerVector implements FieldVe * The index is the type id, and the value is the type field. */ private Field[] typeFields = new Field[Byte.MAX_VALUE + 1]; + + /** + * The count of values for each type-id. + */ + private int[] typeCounts = new int[Byte.MAX_VALUE + 1]; + /** * The index is the index into the typeFields array, and the value is the logical field id. */ @@ -189,6 +195,12 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers offsetBufferAllocationSizeInBytes = offsetBuffer.capacity(); this.valueCount = fieldNode.getLength(); + + Arrays.fill(typeCounts, 0); + for (int i = 0; i < this.valueCount; i++) { + byte typeId = getTypeId(i); + if (typeId >= 0) typeCounts[typeId]++; + } } @Override @@ -236,9 +248,11 @@ public synchronized byte registerNewTypeId(Field field) { 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) { + if (fieldType != null) { + int[] typeIds = ((ArrowType.Union) fieldType.getType()).getTypeIds(); + // HACK (JH) I'm not sure how this works for cases when you _do_ want a manual typeId mapping + // - I suspect some reused type ids :| + if (typeIds != null && nextTypeId < typeIds.length) { int thisTypeId = typeIds[nextTypeId]; if (thisTypeId > Byte.MAX_VALUE) { throw new IllegalStateException("Dense union vector types must be bytes. " + thisTypeId + " is too large"); @@ -447,7 +461,7 @@ private void reallocTypeBuffer() { typeBuffer.getReferenceManager().release(1); typeBuffer = newBuf; typeBufferAllocationSizeInBytes = (int)newAllocationSize; - setNegative(currentBufferCapacity, newBuf.capacity() - currentBufferCapacity); + setNegative(currentBufferCapacity, newBuf.capacity()); } private void reallocOffsetBuffer() { @@ -501,6 +515,7 @@ public void close() { @Override public void clear() { valueCount = 0; + Arrays.fill(typeCounts, 0); typeBuffer.getReferenceManager().release(); typeBuffer = allocator.getEmpty(); offsetBuffer.getReferenceManager().release(); @@ -511,6 +526,7 @@ public void clear() { @Override public void reset() { valueCount = 0; + Arrays.fill(typeCounts, 0); setNegative(0, typeBuffer.capacity()); offsetBuffer.setZero(0, offsetBuffer.capacity()); internalStruct.reset(); @@ -522,7 +538,8 @@ public Field getField() { List childFields = new ArrayList<>(childCount); int[] typeIds = new int[childCount]; for (int i = 0; i < typeFields.length; i++) { - if (typeFields[i] != null) { + if (childVectors[i] != null) { + typeFields[i] = childVectors[i].getField(); int curIdx = childFields.size(); typeIds[curIdx] = i; childFields.add(typeFields[i]); @@ -531,11 +548,12 @@ public Field getField() { FieldType fieldType; if (this.fieldType == null) { - fieldType = FieldType.nullable(new ArrowType.Union(Dense, typeIds)); + fieldType = FieldType.notNullable(new ArrowType.Union(Dense, null)); } else { final UnionMode mode = UnionMode.Dense; - fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds), - this.fieldType.getDictionary(), this.fieldType.getMetadata()); + boolean indirectTypeIds = ((Union) this.fieldType.getType()).getTypeIds() != null; + fieldType = new FieldType(false, new ArrowType.Union(mode, indirectTypeIds ? typeIds : null), + this.fieldType.getDictionary(), this.fieldType.getMetadata()); } return new Field(name, fieldType, childFields); @@ -587,6 +605,18 @@ public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } + FieldVector putVector(byte typeId, FieldVector child) { + internalStruct.putVector(child.getName(), child); + childVectors[typeId] = child; + typeFields[typeId] = child.getField(); + + if (callBack != null) { + callBack.doWork(); + } + + return child; + } + 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)); @@ -623,12 +653,13 @@ public TransferImpl(DenseUnionVector to) { } private void createTransferPairs() { - for (int i = 0; i < nextTypeId; i++) { + for (byte i = 0; i < nextTypeId; i++) { + byte typeId = typeMapFields[i]; ValueVector srcVec = internalStruct.getVectorById(i); ValueVector dstVec = to.internalStruct.getVectorById(i); - to.typeFields[i] = typeFields[i]; + to.typeFields[typeId] = typeFields[typeId]; to.typeMapFields[i] = typeMapFields[i]; - to.childVectors[i] = dstVec; + to.childVectors[typeId] = dstVec; internalTransferPairs[i] = srcVec.makeTransferPair(dstVec); } } @@ -643,13 +674,16 @@ public void transfer() { refManager = offsetBuffer.getReferenceManager(); to.offsetBuffer = refManager.transferOwnership(offsetBuffer, to.allocator).getTransferredBuffer(); - for (int i = 0; i < nextTypeId; i++) { + for (byte i = 0; i < nextTypeId; i++) { + byte typeId = typeMapFields[i]; if (internalTransferPairs[i] != null) { internalTransferPairs[i].transfer(); - to.childVectors[i] = internalTransferPairs[i].getTo(); + to.childVectors[typeId] = internalTransferPairs[i].getTo(); } } + to.typeCounts = typeCounts.clone(); to.valueCount = valueCount; + to.nextTypeId = nextTypeId; clear(); } @@ -665,36 +699,51 @@ public void splitAndTransfer(int startIndex, int length) { to.typeBuffer = refManager.transferOwnership(slicedBuffer, to.allocator).getTransferredBuffer(); // transfer offset buffer - while (to.offsetBuffer.capacity() < (long) length * OFFSET_WIDTH) { - to.reallocOffsetBuffer(); - } + int [] typeCounts; + int [] typeStarts = new int[Byte.MAX_VALUE + 1]; - int [] typeCounts = new int[nextTypeId]; - int [] typeStarts = new int[nextTypeId]; - for (int i = 0; i < typeCounts.length; i++) { - typeCounts[i] = 0; - typeStarts[i] = -1; - } + if (startIndex == 0 && length == valueCount) { + // transferring the whole vector, skip the offset mapping + int[] srcTypeCounts = DenseUnionVector.this.typeCounts; + typeCounts = Arrays.copyOf(srcTypeCounts, srcTypeCounts.length); + + ArrowBuf slicedOffsetBuffer = offsetBuffer.slice(0, (long) length * OFFSET_WIDTH); + ReferenceManager offsetRefManager = slicedOffsetBuffer.getReferenceManager(); + to.offsetBuffer = offsetRefManager.transferOwnership(slicedOffsetBuffer, to.allocator).getTransferredBuffer(); + } else { + typeCounts = new int[Byte.MAX_VALUE + 1]; + + for (int i = 0; i < typeCounts.length; i++) { + 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); + while (to.offsetBuffer.capacity() < (long) length * OFFSET_WIDTH) { + to.reallocOffsetBuffer(); + } + + 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(); + byte typeId = typeMapFields[i]; + if (typeCounts[typeId] > 0 && typeStarts[typeId] != -1) { + internalTransferPairs[i].splitAndTransfer(typeStarts[typeId], typeCounts[typeId]); } - } + to.childVectors[typeId] = internalTransferPairs[i].getTo(); + } + to.typeCounts = typeCounts; + to.nextTypeId = nextTypeId; to.setValueCount(length); } @@ -830,19 +879,9 @@ public void setValueCount(int valueCount) { 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]]); + childVectors[typeMapFields[i]].setValueCount(typeCounts[typeMapFields[i]]); } } @@ -905,7 +944,11 @@ public void setTypeId(int index, byte typeId) { while (index >= getTypeBufferValueCapacity()) { reallocTypeBuffer(); } + int oldTypeId = getTypeId(index); + if (oldTypeId >= 0) typeCounts[oldTypeId]--; + typeBuffer.setByte(index * TYPE_WIDTH , typeId); + if (typeId >= 0) typeCounts[typeId]++; } private int getTypeBufferValueCapacity() { diff --git a/vector/src/main/codegen/templates/UnionVector.java b/vector/src/main/codegen/templates/UnionVector.java index 67efdf60f7..8e403d2db7 100644 --- a/vector/src/main/codegen/templates/UnionVector.java +++ b/vector/src/main/codegen/templates/UnionVector.java @@ -848,10 +848,14 @@ public void setSafe(int index, Nullable${name}Holder holder) { public void setType(int index, MinorType type) { + setTypeId(index, (byte) type.ordinal()); + } + + public void setTypeId(int index, byte typeId) { while (index >= getTypeBufferValueCapacity()) { reallocTypeBuffer(); } - typeBuffer.setByte(index * TYPE_WIDTH , (byte) type.ordinal()); + typeBuffer.setByte(index * TYPE_WIDTH , typeId); } private int getTypeBufferValueCapacity() { diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java index 68700fe6af..1080b4988b 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java @@ -45,6 +45,7 @@ import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; +import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.Float4Vector; @@ -122,6 +123,7 @@ public JsonFileWriter(File outputFile) throws IOException { } /** Constructs a new writer that will output to outputFile with the given options. */ + @SuppressWarnings("deprecation") public JsonFileWriter(File outputFile, JSONWriteConfig config) throws IOException { MappingJsonFactory jsonFactory = new MappingJsonFactory(); this.generator = jsonFactory.createGenerator(outputFile, JsonEncoding.UTF8); @@ -202,6 +204,11 @@ private void writeBatch(VectorSchemaRoot recordBatch) throws IOException { } private void writeFromVectorIntoJson(Field field, FieldVector vector) throws IOException { + if (vector instanceof ExtensionTypeVector) { + writeFromVectorIntoJson(field, ((ExtensionTypeVector) vector).getUnderlyingVector()); + return; + } + TypeLayout typeLayout = TypeLayout.getTypeLayout(field.getType()); List vectorTypes = typeLayout.getBufferTypes(); List vectorBuffers = vector.getFieldBuffers(); @@ -424,6 +431,7 @@ private void writeValueToGenerator( case VARBINARY: case LIST: case MAP: + case DENSEUNION: generator.writeNumber(buffer.getInt((long) index * BaseVariableWidthVector.OFFSET_WIDTH)); break; case LISTVIEW: