/** * Load a record batch from a single buffer. * * @param def The definition for the record batch. * @param buf The buffer that holds the data associated with the record batch * @return Whether or not the schema changed since the previous load. * @throws SchemaChangeException */ public boolean load(RecordBatchDef def, ByteBuf buf) throws SchemaChangeException { // logger.debug("Loading record batch with def {} and data {}", def, buf); this.valueCount = def.getRecordCount(); boolean schemaChanged = schema == null; Map<MaterializedField, ValueVector> oldFields = Maps.newHashMap(); for (VectorWrapper<?> w : container) { ValueVector v = w.getValueVector(); oldFields.put(v.getField(), v); } VectorContainer newVectors = new VectorContainer(); List<FieldMetadata> fields = def.getFieldList(); int bufOffset = 0; for (FieldMetadata fmd : fields) { FieldDef fieldDef = fmd.getDef(); ValueVector v = oldFields.remove(fieldDef); if (v != null) { container.add(v); continue; } // if we arrive here, we didn't have a matching vector. schemaChanged = true; MaterializedField m = new MaterializedField(fieldDef); v = TypeHelper.getNewVector(m, allocator); if (fmd.getValueCount() == 0) { v.clear(); } else { v.load(fmd, buf.slice(bufOffset, fmd.getBufferLength())); } bufOffset += fmd.getBufferLength(); newVectors.add(v); } if (!oldFields.isEmpty()) { schemaChanged = true; for (ValueVector v : oldFields.values()) { v.close(); } } // rebuild the schema. SchemaBuilder b = BatchSchema.newBuilder(); for (VectorWrapper<?> v : newVectors) { b.addField(v.getField()); } b.setSelectionVectorMode(BatchSchema.SelectionVectorMode.NONE); this.schema = b.build(); container = newVectors; return schemaChanged; }
/** * 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; }
public VectorWrapper<?> getValueAccessorById(int fieldId, Class<?> clazz) { return container.getVectorAccessor(fieldId, clazz); }
public TypedFieldId getValueVectorId(SchemaPath path) { return container.getValueVector(path); }