@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(); }
@Override public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) { count.addAndGet(result.getHeader().getRowCount()); child.dataArrived(result, throttle); }
@Override public void dataArrived(final QueryDataBatch result, final ConnectionThrottle throttle) { final QueryData queryData = result.getHeader(); if (result.hasData()) { final int nRows = this.nRows.addAndGet(queryData.getRowCount()); if (nRows > 1) { throw new IllegalStateException("Expected exactly one row, but got " + nRows); } rowArrived(result); } result.release(); }
@Test public void twoBitTwoExchange() throws Exception { try (final ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); final SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); final SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); final DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator);) { bit1.run(); bit2.run(); client.connect(); final List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/mergerecv/merging_receiver.json"), Charsets.UTF_8)); int count = 0; final RecordBatchLoader batchLoader = new RecordBatchLoader(client.getRecordAllocator()); // print the results for (final QueryDataBatch b : results) { final QueryData queryData = b.getHeader(); final int rowCount = queryData.getRowCount(); count += rowCount; batchLoader.load(queryData.getDef(), b.getData()); // loaded but not used, just to test b.release(); batchLoader.clear(); } assertEquals(200000, count); } }
@Override public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) { QueryData queryHeader = result.getHeader(); int rows = queryHeader.getRowCount(); try { if ( result.hasData() ) { ArrowBuf data = result.getData(); loader.load(queryHeader.getDef(), data); for (int i = 0; i < rows; i++) { Map<String,String> record = Maps.newHashMap(); for (VectorWrapper<?> vw : loader) { final String field = vw.getValueVector().getField().getName(); final ValueVector vv = vw.getValueVector(); final Object value = i < vv.getValueCount() ? vv.getObject(i) : null; final String display = value == null ? null : value.toString(); record.put(field, display); } records.add(record); } loader.clear(); } result.release(); } catch (SchemaChangeException e) { fail(e.getMessage()); } }
int count = 0; for (QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); if (b.getHeader().getRowCount() == 0) { break;
@Ignore("DX-3872") @Test public void testFullExecution() throws Exception { try(ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { bit1.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString( FileUtils.getResourceAsFile("/store/text/test.json"), Charsets.UTF_8) .replace("#{DATA_FILE}", FileUtils.getResourceAsFile("/store/text/data/regions.csv").toURI().toString())); int count = 0; RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator()); for(QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } loader.load(b.getHeader().getDef(), b.getData()); VectorUtil.showVectorAccessibleContent(loader); loader.clear(); b.release(); } assertEquals(5, count); } }
@Ignore("DX-3872") @Test public void TestSingleBroadcastExchangeWithTwoScans() throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { bit1.run(); bit2.run(); client.connect(); String physicalPlan = Files.toString( FileUtils.getResourceAsFile("/sender/broadcast_exchange.json"), Charsets.UTF_8) .replace("#{LEFT_FILE}", FileUtils.getResourceAsFile("/join/merge_single_batch.left.json").toURI().toString()) .replace("#{RIGHT_FILE}", FileUtils.getResourceAsFile("/join/merge_single_batch.right.json").toURI().toString()); List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan); int count = 0; for (QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } b.release(); } assertEquals(25, count); } }
@Ignore("DX-3872") @Test public void multipleSendLocationBroadcastExchange() throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { bit1.run(); bit2.run(); client.connect(); String physicalPlan = Files.toString( FileUtils.getResourceAsFile("/sender/broadcast_exchange_long_run.json"), Charsets.UTF_8); List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan); int count = 0; for (QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } b.release(); } System.out.println(count); } }
@Test public void twoBitTwoExchangeTwoEntryRun() throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { bit1.run(); bit2.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/sender/hash_exchange.json"), Charsets.UTF_8)); int count = 0; for (QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } b.release(); } assertEquals(200, count); } }
@Test public void twoBitTwoExchangeTwoEntryRun() throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { bit1.run(); bit2.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/sender/union_exchange.json"), Charsets.UTF_8)); int count = 0; for (QueryDataBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } b.release(); } assertEquals(150, count); } }
private List<List<String>> getOutput(List<QueryDataBatch> batches) throws SchemaChangeException { List<List<String>> output = new ArrayList<>(); RecordBatchLoader loader = new RecordBatchLoader(getAllocator()); int last = 0; for(QueryDataBatch batch : batches) { int rows = batch.getHeader().getRowCount(); if(batch.getData() != null) { loader.load(batch.getHeader().getDef(), batch.getData()); // TODO: Clean: DRILL-2933: That load(...) no longer throws // SchemaChangeException, so check/clean throws clause above. for (int i = 0; i < rows; ++i) { output.add(new ArrayList<String>()); for (VectorWrapper<?> vw: loader) { ValueVector vv = vw.getValueVector(); Object o = vv.getObject(i); output.get(last).add(o == null ? null: o.toString()); } ++last; } } loader.clear(); batch.release(); } return output; }
@Test public void twoBitOneExchangeTwoEntryRun() throws Exception{ try(ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); SabotNode bit2 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, false); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)){ bit1.run(); bit2.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/physical_single_exchange_double_entry.json"), Charsets.UTF_8)); int count = 0; for(QueryDataBatch b : results){ count += b.getHeader().getRowCount(); b.release(); } assertEquals(200, count); } } }
@Test public void oneBitOneExchangeOneEntryRun() throws Exception{ try(ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)){ bit1.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/physical_single_exchange.json"), Charsets.UTF_8)); int count = 0; for(QueryDataBatch b : results){ count += b.getHeader().getRowCount(); b.release(); } assertEquals(100, count); } }
@Test public void oneBitOneExchangeTwoEntryRun() throws Exception{ try(ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit1 = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)){ bit1.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/physical_single_exchange_double_entry.json"), Charsets.UTF_8)); int count = 0; for(QueryDataBatch b : results){ count += b.getHeader().getRowCount(); b.release(); } assertEquals(200, count); } }
protected int printResult(List<QueryDataBatch> results) throws SchemaChangeException { int rowCount = 0; final RecordBatchLoader loader = new RecordBatchLoader(getAllocator()); for(final QueryDataBatch result : results) { rowCount += result.getHeader().getRowCount(); loader.load(result.getHeader().getDef(), result.getData()); // TODO: Clean: DRILL-2933: That load(...) no longer throws // SchemaChangeException, so check/clean throw clause above. VectorUtil.showVectorAccessibleContent(loader, columnWidths); loader.clear(); result.release(); } System.out.println("Total record count: " + rowCount); return rowCount; }
public Builder mergeFrom(com.dremio.exec.proto.UserBitShared.QueryData other) { if (other == com.dremio.exec.proto.UserBitShared.QueryData.getDefaultInstance()) return this; if (other.hasQueryId()) { mergeQueryId(other.getQueryId()); } if (other.hasRowCount()) { setRowCount(other.getRowCount()); } if (other.hasDef()) { mergeDef(other.getDef()); } this.mergeUnknownFields(other.getUnknownFields()); return this; }
public static int getRecordCount(List<QueryDataBatch> result) throws UnsupportedEncodingException, SchemaChangeException { int i = 0; for (QueryDataBatch batch : result) { i += batch.getHeader().getRowCount(); batch.release(); } return i; }
@Override public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) { final int rows = result.getHeader().getRowCount(); if (result.getData() != null) { count.addAndGet(rows); } result.release(); }
public void writeTo(io.protostuff.Output output, com.dremio.exec.proto.UserBitShared.QueryData message) throws java.io.IOException { if(message.hasQueryId()) output.writeObject(1, message.getQueryId(), com.dremio.exec.proto.SchemaUserBitShared.QueryId.WRITE, false); if(message.hasRowCount()) output.writeInt32(2, message.getRowCount(), false); if(message.hasDef()) output.writeObject(3, message.getDef(), com.dremio.exec.proto.SchemaUserBitShared.RecordBatchDef.WRITE, false); } public boolean isInitialized(com.dremio.exec.proto.UserBitShared.QueryData message)