public void testCommon(String[] expectedResults, String physicalPlan, String resourceFile)
      throws Exception {
    try (RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        Drillbit bit = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {

      // run query.
      bit.run();
      client.connect();
      List<QueryDataBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Files.toString(FileUtils.getResourceAsFile(physicalPlan), Charsets.UTF_8)
                  .replace("#{TEST_FILE}", resourceFile));

      RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());

      QueryDataBatch batch = results.get(0);
      assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));

      int i = 0;
      for (VectorWrapper<?> v : batchLoader) {

        ValueVector.Accessor accessor = v.getValueVector().getAccessor();
        System.out.println(accessor.getObject(0));
        assertEquals(expectedResults[i++], accessor.getObject(0).toString());
      }

      batchLoader.clear();
      for (QueryDataBatch b : results) {
        b.release();
      }
    }
  }
Esempio n. 2
0
 @Override
 public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) {
   try {
     final int rows = result.getHeader().getRowCount();
     if (result.hasData()) {
       RecordBatchLoader loader = null;
       try {
         loader = new RecordBatchLoader(allocator);
         loader.load(result.getHeader().getDef(), result.getData());
         // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
         // SchemaChangeException, so check/clean catch clause below.
         for (int i = 0; i < loader.getSchema().getFieldCount(); ++i) {
           columns.add(loader.getSchema().getColumn(i).getPath());
         }
         for (int i = 0; i < rows; ++i) {
           final Map<String, String> record = Maps.newHashMap();
           for (VectorWrapper<?> vw : loader) {
             final String field = vw.getValueVector().getMetadata().getNamePart().getName();
             final ValueVector.Accessor accessor = vw.getValueVector().getAccessor();
             final Object value = i < accessor.getValueCount() ? accessor.getObject(i) : null;
             final String display = value == null ? null : value.toString();
             record.put(field, display);
           }
           results.add(record);
         }
       } finally {
         if (loader != null) {
           loader.clear();
         }
       }
     }
   } catch (SchemaChangeException e) {
     throw new RuntimeException(e);
   } finally {
     result.release();
   }
 }