private static Map<String, MajorType> extractSchema(final RecordBatchDef batchDef) { final List<SerializedField> serializedFields = batchDef.getFieldList(); final Map<String, MajorType> fieldsMap = Maps.newHashMap(); for (SerializedField field : serializedFields) { fieldsMap.put(field.getNamePart().getName(), field.getMajorType()); } return fieldsMap; } }
/** * computes the number of buffers for a given serialized field * @param field serialized field * @param buffers total buffers in the batch * @param buffersStart starting buffer for the passed field * * @return number of buffers for the field */ private static int fieldBuffersCount(SerializedField field, ByteBuf[] buffers, final int buffersStart) { int totalBufferWidth = 0; int lastIndex = buffersStart; while (totalBufferWidth < field.getBufferLength() && lastIndex < buffersStart + buffers.length) { ByteBuf buf = buffers[lastIndex]; totalBufferWidth += buf.readableBytes(); ++lastIndex; } if (totalBufferWidth != field.getBufferLength()) { throw new IllegalStateException("not enough buffers for field " + field.getNamePart().getName() + " of type " + field.getMajorType()); } return lastIndex - buffersStart; }
public static Field getFieldForSerializedField(SerializedField serializedField) { String name = serializedField.getNamePart().getName(); org.apache.arrow.vector.types.Types.MinorType arrowMinorType = getArrowMinorType(serializedField.getMajorType().getMinorType()); switch(serializedField.getMajorType().getMinorType()) { case LIST: return new Field(name, true, arrowMinorType.getType(), ImmutableList.of(getFieldForSerializedField(serializedField.getChild(2)))); return new Field(name, true, new Decimal(serializedField.getMajorType().getPrecision(), serializedField.getMajorType().getScale()), null); default: return new Field(name, true, arrowMinorType.getType(), null);
public Builder mergeFrom(com.dremio.exec.proto.UserBitShared.SerializedField other) { if (other == com.dremio.exec.proto.UserBitShared.SerializedField.getDefaultInstance()) return this; if (other.hasMajorType()) { mergeMajorType(other.getMajorType());
Assert.assertEquals(MinorType.VARCHAR, results.get(0).getHeader().getDef().getField(1).getChild(2).getMajorType().getMinorType()); } finally { for (QueryDataBatch batch : results) {
public void load(SerializedField metadata, ArrowBuf buffer) { /* release the current buffers (if any) */ listVector.clear(); /* load inner offset buffer */ final SerializedField offsetMetadata = metadata.getChild(0); final int offsetLength = offsetMetadata.getBufferLength(); loadOffsetBuffer(offsetMetadata, buffer); /* load inner validity buffer */ final SerializedField bitMetadata = metadata.getChild(1); final int bitLength = bitMetadata.getBufferLength(); loadValidityBuffer(bitMetadata, buffer.slice(offsetLength, bitLength)); /* load inner data vector */ final SerializedField vectorMetadata = metadata.getChild(2); if (listVector.getDataVector() == BaseRepeatedValueVector.DEFAULT_DATA_VECTOR) { listVector.addOrGetVector(FieldType.nullable(getArrowMinorType(vectorMetadata.getMajorType().getMinorType()).getType())); } final int vectorLength = vectorMetadata.getBufferLength(); TypeHelper.load(listVector.vector, vectorMetadata, buffer.slice(offsetLength + bitLength, vectorLength)); listVector.setLastSet(metadata.getValueCount()); listVector.valueCount = metadata.getValueCount(); }
public void writeTo(io.protostuff.Output output, com.dremio.exec.proto.UserBitShared.SerializedField message) throws java.io.IOException { if(message.hasMajorType()) output.writeObject(1, message.getMajorType(), com.dremio.common.types.SchemaTypeProtos.MajorType.WRITE, false); if(message.hasNamePart()) output.writeObject(2, message.getNamePart(), com.dremio.exec.proto.SchemaUserBitShared.NamePart.WRITE, false); for(com.dremio.exec.proto.UserBitShared.SerializedField child : message.getChildList()) output.writeObject(3, child, com.dremio.exec.proto.SchemaUserBitShared.SerializedField.WRITE, true); if(message.hasValueCount()) output.writeInt32(4, message.getValueCount(), false); if(message.hasVarByteLength()) output.writeInt32(5, message.getVarByteLength(), false); if(message.hasBufferLength()) output.writeInt32(7, message.getBufferLength(), false); } public boolean isInitialized(com.dremio.exec.proto.UserBitShared.SerializedField message)
public void load(SerializedField metadata, ArrowBuf buffer) { /* release the current buffers (if any) */ vector.clear(); /* get the metadata for all children */ final SerializedField offsetMetadata = metadata.getChild(0); final SerializedField vectorMetadata = metadata.getChild(1); final int offsetLength = offsetMetadata.getBufferLength(); final int vectorLength = vectorMetadata.getBufferLength(); /* load inner offset buffer */ loadOffsetBuffer(offsetMetadata, buffer); /* load inner data vector */ if (vector.getDataVector() == BaseRepeatedValueVector.DEFAULT_DATA_VECTOR) { vector.addOrGetVector(FieldType.nullable(getArrowMinorType(metadata.getMajorType().getMinorType()).getType())); } TypeHelper.load(vector.vector, vectorMetadata, buffer.slice(offsetLength, vectorLength)); }
private static boolean typeEqual(SerializedField field1, SerializedField field2) { if (!field1.getMajorType().equals(field2.getMajorType())) { return false; } if (field1.getChildCount() != field2.getChildCount()) { return false; } for(int i=0; i<field1.getChildCount(); i++) { final SerializedField childField1 = field1.getChild(i); final SerializedField childField2 = field2.getChild(i); if (!typeEqual(childField1, childField2)) { return false; } } return true; }
@Test public void test() throws Exception { ElasticsearchCluster.ColumnData[] data = new ElasticsearchCluster.ColumnData[]{ new ElasticsearchCluster.ColumnData("column", TEXT, new Object[][]{ {"value"} }) }; elastic.load(schema, table, data); List<QueryDataBatch> results = testSqlWithResults(String.format("select * from elasticsearch.%s.%s limit 0", schema, table)); MajorType t = results.get(0).getHeader().getDef().getField(0).getMajorType(); Assert.assertEquals("Mismatched type", MinorType.VARCHAR, t.getMinorType()); Assert.assertEquals("Mismatched datamode", DataMode.OPTIONAL, t.getMode()); } }
private void checkResult(RecordBatchLoader loader, List<QueryDataBatch> results) { loader.load(results.get(0).getHeader().getDef(), results.get(0).getData()); RecordBatchDef def = results.get(0).getHeader().getDef(); // the entire row is returned as a single column assertEquals(1, def.getFieldCount()); assertTrue(def.getField(0).getMajorType().getMode() == DataMode.OPTIONAL); loader.clear(); for(QueryDataBatch result : results) { result.release(); } client.close(); }
private Map<SchemaPath, MajorType> getTypeMapFromBatch(QueryDataBatch batch) { Map<SchemaPath, MajorType> typeMap = new TreeMap<>(); for (int i = 0; i < batch.getHeader().getDef().getFieldCount(); i++) { typeMap.put(SchemaPath.getSimplePath(SerializedFieldHelper.create(batch.getHeader().getDef().getField(i)).getName()), batch.getHeader().getDef().getField(i).getMajorType()); } return typeMap; }
@Test public void testGetType() throws Exception { assertEquals(UserBitShared.SerializedField.getDefaultInstance().getMajorType(), genericAccessor.getType()); } }