@Test public void TestMultipleSendLocationBroadcastExchange() throws Exception { RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet); Drillbit bit2 = new Drillbit(CONFIG, serviceSet); DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) { bit1.run(); bit2.run(); client.connect(); String physicalPlan = Files.toString( FileUtils.getResourceAsFile("/sender/broadcast_exchange_long_run.json"), Charsets.UTF_8); List<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan); int count = 0; for (QueryResultBatch b : results) { if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount(); b.release(); } System.out.println(count); } }
@Test public void twoBitTwoExchangeTwoEntryRun() throws Exception { RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet); Drillbit bit2 = new Drillbit(CONFIG, serviceSet); DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) { bit1.run(); bit2.run(); client.connect(); List<QueryResultBatch> results = client.runQuery( org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString( FileUtils.getResourceAsFile("/sender/union_exchange.json"), Charsets.UTF_8)); int count = 0; for (QueryResultBatch b : results) { if (b.getHeader().getRowCount() != 0) { count += b.getHeader().getRowCount(); } b.release(); } assertEquals(150, count); } }
@Test public void TestSingleBroadcastExchangeWithTwoScans() throws Exception { RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet); Drillbit bit2 = new Drillbit(CONFIG, serviceSet); DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) { 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<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan); int count = 0; for (QueryResultBatch b : results) { if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount(); b.release(); } assertEquals(25, count); } }
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(); } } }
@Test public void testMultipleProvidersMixedSizes() throws Exception { RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet); Drillbit bit2 = new Drillbit(CONFIG, serviceSet); DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) { bit1.run(); bit2.run(); client.connect(); List<QueryResultBatch> results = client.runQuery( org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString( FileUtils.getResourceAsFile("/mergerecv/multiple_providers.json"), Charsets.UTF_8)); int count = 0; RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator()); // print the results Long lastBlueValue = null; for (QueryResultBatch b : results) { count += b.getHeader().getRowCount(); for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) { List<Object> row = Lists.newArrayList(); batchLoader.load(b.getHeader().getDef(), b.getData()); for (VectorWrapper vw : batchLoader) { row.add( vw.getValueVector().getField().toExpr() + ":" + vw.getValueVector().getAccessor().getObject(valueIdx)); if (vw.getValueVector() .getField() .getAsSchemaPath() .getRootSegment() .getPath() .equals("blue")) { // assert order is ascending if (((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue() == 0) continue; // ignore initial 0's from sort if (lastBlueValue != null) assertTrue( ((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue() >= ((Long) lastBlueValue).longValue()); lastBlueValue = (Long) vw.getValueVector().getAccessor().getObject(valueIdx); } } for (Object cell : row) { int len = cell.toString().length(); System.out.print(cell + " "); for (int i = 0; i < (30 - len); ++i) System.out.print(" "); } System.out.println(); } b.release(); batchLoader.clear(); } assertEquals(400, count); } }
@Test public void twoBitTwoExchange() throws Exception { RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet); Drillbit bit2 = new Drillbit(CONFIG, serviceSet); DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) { bit1.run(); bit2.run(); client.connect(); List<QueryResultBatch> results = client.runQuery( org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString( FileUtils.getResourceAsFile("/mergerecv/merging_receiver.json"), Charsets.UTF_8)); int count = 0; RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator()); // print the results for (QueryResultBatch b : results) { count += b.getHeader().getRowCount(); for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) { List<Object> row = Lists.newArrayList(); batchLoader.load(b.getHeader().getDef(), b.getData()); for (VectorWrapper<?> vw : batchLoader) row.add( vw.getValueVector().getField().toExpr() + ":" + vw.getValueVector().getAccessor().getObject(valueIdx)); for (Object cell : row) { if (cell == null) { System.out.print("<null> "); continue; } int len = cell.toString().length(); System.out.print(cell + " "); for (int i = 0; i < (30 - len); ++i) System.out.print(" "); } System.out.println(); } b.release(); batchLoader.clear(); } assertEquals(200, count); } }
@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)); } }
@Test public void testAllocators() throws Exception { // Setup a drillbit (initializes a root allocator) final DrillConfig config = DrillConfig.create(TEST_CONFIGURATIONS); final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet(); final Drillbit bit = new Drillbit(config, serviceSet); bit.run(); final DrillbitContext bitContext = bit.getContext(); FunctionImplementationRegistry functionRegistry = bitContext.getFunctionImplementationRegistry(); StoragePluginRegistry storageRegistry = new StoragePluginRegistry(bitContext); // Create a few Fragment Contexts BitControl.PlanFragment.Builder pfBuilder1 = BitControl.PlanFragment.newBuilder(); pfBuilder1.setMemInitial(1500000); BitControl.PlanFragment pf1 = pfBuilder1.build(); BitControl.PlanFragment.Builder pfBuilder2 = BitControl.PlanFragment.newBuilder(); pfBuilder2.setMemInitial(500000); BitControl.PlanFragment pf2 = pfBuilder1.build(); FragmentContext fragmentContext1 = new FragmentContext(bitContext, pf1, null, functionRegistry); FragmentContext fragmentContext2 = new FragmentContext(bitContext, pf2, null, functionRegistry); // Get a few physical operators. Easiest way is to read a physical plan. PhysicalPlanReader planReader = new PhysicalPlanReader( config, config.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance(), storageRegistry); PhysicalPlan plan = planReader.readPhysicalPlan( Files.toString(FileUtils.getResourceAsFile(planFile), Charsets.UTF_8)); List<PhysicalOperator> physicalOperators = plan.getSortedOperators(); Iterator<PhysicalOperator> physicalOperatorIterator = physicalOperators.iterator(); PhysicalOperator physicalOperator1 = physicalOperatorIterator.next(); PhysicalOperator physicalOperator2 = physicalOperatorIterator.next(); PhysicalOperator physicalOperator3 = physicalOperatorIterator.next(); PhysicalOperator physicalOperator4 = physicalOperatorIterator.next(); PhysicalOperator physicalOperator5 = physicalOperatorIterator.next(); PhysicalOperator physicalOperator6 = physicalOperatorIterator.next(); // Create some bogus Operator profile defs and stats to create operator contexts OpProfileDef def; OperatorStats stats; // Use some bogus operator type to create a new operator context. def = new OpProfileDef( physicalOperator1.getOperatorId(), UserBitShared.CoreOperatorType.MOCK_SUB_SCAN_VALUE, OperatorContext.getChildCount(physicalOperator1)); stats = fragmentContext1.getStats().getOperatorStats(def, fragmentContext1.getAllocator()); // Add a couple of Operator Contexts // Initial allocation = 1000000 bytes for all operators OperatorContext oContext11 = fragmentContext1.newOperatorContext(physicalOperator1, true); DrillBuf b11 = oContext11.getAllocator().buffer(1000000); OperatorContext oContext12 = fragmentContext1.newOperatorContext(physicalOperator2, stats, true); DrillBuf b12 = oContext12.getAllocator().buffer(500000); OperatorContext oContext21 = fragmentContext1.newOperatorContext(physicalOperator3, true); def = new OpProfileDef( physicalOperator4.getOperatorId(), UserBitShared.CoreOperatorType.TEXT_WRITER_VALUE, OperatorContext.getChildCount(physicalOperator4)); stats = fragmentContext2.getStats().getOperatorStats(def, fragmentContext2.getAllocator()); OperatorContext oContext22 = fragmentContext2.newOperatorContext(physicalOperator4, stats, true); DrillBuf b22 = oContext22.getAllocator().buffer(2000000); // New Fragment begins BitControl.PlanFragment.Builder pfBuilder3 = BitControl.PlanFragment.newBuilder(); pfBuilder3.setMemInitial(1000000); BitControl.PlanFragment pf3 = pfBuilder3.build(); FragmentContext fragmentContext3 = new FragmentContext(bitContext, pf3, null, functionRegistry); // New fragment starts an operator that allocates an amount within the limit def = new OpProfileDef( physicalOperator5.getOperatorId(), UserBitShared.CoreOperatorType.UNION_VALUE, OperatorContext.getChildCount(physicalOperator5)); stats = fragmentContext3.getStats().getOperatorStats(def, fragmentContext3.getAllocator()); OperatorContext oContext31 = fragmentContext3.newOperatorContext(physicalOperator5, stats, true); DrillBuf b31a = oContext31.getAllocator().buffer(200000); // Previously running operator completes b22.release(); ((AutoCloseable) oContext22).close(); // Fragment 3 asks for more and fails boolean outOfMem = false; try { DrillBuf b31b = oContext31.getAllocator().buffer(4400000); if (b31b != null) { b31b.release(); } else { outOfMem = true; } } catch (Exception e) { outOfMem = true; } assertEquals(true, (boolean) outOfMem); // Operator is Exempt from Fragment limits. Fragment 3 asks for more and succeeds outOfMem = false; OperatorContext oContext32 = fragmentContext3.newOperatorContext(physicalOperator6, false); DrillBuf b32 = null; try { b32 = oContext32.getAllocator().buffer(4400000); } catch (Exception e) { outOfMem = true; } finally { if (b32 != null) { b32.release(); } else { outOfMem = true; } closeOp(oContext32); } assertEquals(false, (boolean) outOfMem); b11.release(); closeOp(oContext11); b12.release(); closeOp(oContext12); closeOp(oContext21); b31a.release(); closeOp(oContext31); fragmentContext1.close(); fragmentContext2.close(); fragmentContext3.close(); bit.close(); serviceSet.close(); }