/** * Creates a copy a record batch, converting any fields as necessary to coerce it into the * provided schema * * @param in * @param toSchema * @param context * @return */ public static VectorContainer coerceContainer( VectorAccessible in, BatchSchema toSchema, OperatorContext context) { int recordCount = in.getRecordCount(); Map<SchemaPath, ValueVector> vectorMap = Maps.newHashMap(); for (VectorWrapper w : in) { ValueVector v = w.getValueVector(); vectorMap.put(v.getField().getPath(), v); } VectorContainer c = new VectorContainer(context); for (MaterializedField field : toSchema) { ValueVector v = vectorMap.remove(field.getPath()); if (v != null) { int valueCount = v.getAccessor().getValueCount(); TransferPair tp = v.getTransferPair(); tp.transfer(); if (v.getField().getType().getMinorType().equals(field.getType().getMinorType())) { if (field.getType().getMinorType() == MinorType.UNION) { UnionVector u = (UnionVector) tp.getTo(); for (MinorType t : field.getType().getSubTypeList()) { if (u.getField().getType().getSubTypeList().contains(t)) { continue; } u.addSubType(t); } } c.add(tp.getTo()); } else { ValueVector newVector = TypeHelper.getNewVector(field, context.getAllocator()); Preconditions.checkState( field.getType().getMinorType() == MinorType.UNION, "Can only convert vector to Union vector"); UnionVector u = (UnionVector) newVector; u.addVector(tp.getTo()); MinorType type = v.getField().getType().getMinorType(); for (int i = 0; i < valueCount; i++) { u.getMutator().setType(i, type); } for (MinorType t : field.getType().getSubTypeList()) { if (u.getField().getType().getSubTypeList().contains(t)) { continue; } u.addSubType(t); } u.getMutator().setValueCount(valueCount); c.add(u); } } else { v = TypeHelper.getNewVector(field, context.getAllocator()); v.allocateNew(); v.getMutator().setValueCount(recordCount); c.add(v); } } c.buildSchema(in.getSchema().getSelectionVectorMode()); c.setRecordCount(recordCount); Preconditions.checkState(vectorMap.size() == 0, "Leftover vector from incoming batch"); return c; }