@Override
 public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) {
   int rows = result.getHeader().getRowCount();
   System.out.println(String.format("Result batch arrived. Number of records: %d", rows));
   count.addAndGet(rows);
   result.release();
 }
  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();
      }
    }
  }
  @Override
  public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) {
    final QueryData header = result.getHeader();
    final DrillBuf data = result.getData();

    if (data != null) {
      count.addAndGet(header.getRowCount());
      try {
        loader.load(header.getDef(), data);
        // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
        // SchemaChangeException, so check/clean catch clause below.
      } catch (SchemaChangeException e) {
        submissionFailed(UserException.systemError(e).build(logger));
      }

      switch (format) {
        case TABLE:
          VectorUtil.showVectorAccessibleContent(loader, columnWidth);
          break;
        case TSV:
          VectorUtil.showVectorAccessibleContent(loader, "\t");
          break;
        case CSV:
          VectorUtil.showVectorAccessibleContent(loader, ",");
          break;
      }
      loader.clear();
    }

    result.release();
  }
 public static void setSessionOption(
     final DrillClient drillClient, final String option, final String value) {
   try {
     final List<QueryDataBatch> results =
         drillClient.runQuery(
             UserBitShared.QueryType.SQL,
             String.format("ALTER session SET `%s` = %s", option, value));
     for (final QueryDataBatch data : results) {
       data.release();
     }
   } catch (final RpcException e) {
     fail("Could not set option: " + e.toString());
   }
 }
  @Test
  @Ignore
  public void testParseParquetPhysicalPlan() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
    DrillConfig config = DrillConfig.create();

    try (Drillbit bit1 = new Drillbit(config, serviceSet);
        DrillClient client = new DrillClient(config, serviceSet.getCoordinator()); ) {
      bit1.run();
      client.connect();
      List<QueryDataBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Resources.toString(Resources.getResource(fileName), Charsets.UTF_8));
      RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator());
      int count = 0;
      for (QueryDataBatch b : results) {
        System.out.println(String.format("Got %d results", b.getHeader().getRowCount()));
        count += b.getHeader().getRowCount();
        loader.load(b.getHeader().getDef(), b.getData());
        for (VectorWrapper vw : loader) {
          System.out.print(vw.getValueVector().getField().toExpr() + ": ");
          ValueVector vv = vw.getValueVector();
          for (int i = 0; i < vv.getAccessor().getValueCount(); i++) {
            Object o = vv.getAccessor().getObject(i);
            if (o instanceof byte[]) {
              System.out.print(" [" + new String((byte[]) o) + "]");
            } else {
              System.out.print(" [" + vv.getAccessor().getObject(i) + "]");
            }
            //            break;
          }
          System.out.println();
        }
        loader.clear();
        b.release();
      }
      client.close();
      System.out.println(String.format("Got %d total results", count));
    }
  }
Beispiel #6
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();
   }
 }