Examples with Drillbit org.apache.drill.exec.server.Drillbit used on opensource projects

Example 6 with Drillbit

use of org.apache.drill.exec.server.Drillbit in project drill by axbaretto.

the class QuerySubmitter method submitQuery.

public int submitQuery(String planLocation, String queryString, String type, String zkQuorum, boolean local, int bits, String format, int width) throws Exception {
    DrillConfig config = DrillConfig.create();
    DrillClient client = null;
    Preconditions.checkArgument(!(planLocation == null && queryString == null), "Must provide either query file or query string");
    Preconditions.checkArgument(!(planLocation != null && queryString != null), "Must provide either query file or query string, not both");
    RemoteServiceSet serviceSet = null;
    Drillbit[] drillbits = null;
    try {
        if (local) {
            serviceSet = RemoteServiceSet.getLocalServiceSet();
            drillbits = new Drillbit[bits];
            for (int i = 0; i < bits; i++) {
                drillbits[i] = new Drillbit(config, serviceSet);
                drillbits[i].run();
            }
            client = new DrillClient(config, serviceSet.getCoordinator());
        } else {
            ZKClusterCoordinator clusterCoordinator = new ZKClusterCoordinator(config, zkQuorum);
            clusterCoordinator.start(10000);
            client = new DrillClient(config, clusterCoordinator);
        }
        client.connect();
        String plan;
        if (queryString == null) {
            plan = Charsets.UTF_8.decode(ByteBuffer.wrap(Files.readAllBytes(Paths.get(planLocation)))).toString();
        } else {
            plan = queryString;
        }
        return submitQuery(client, plan, type, format, width);
    } catch (Throwable th) {
        System.err.println("Query Failed due to : " + th.getMessage());
        return -1;
    } finally {
        if (client != null) {
            client.close();
        }
        if (local) {
            for (Drillbit b : drillbits) {
                b.close();
            }
            serviceSet.close();
        }
    }
}

Example 7 with Drillbit

use of org.apache.drill.exec.server.Drillbit in project drill by axbaretto.

the class TestWriteToDisk method test.

@Test
@SuppressWarnings("static-method")
public void test() throws Exception {
    final List<ValueVector> vectorList = Lists.newArrayList();
    final DrillConfig config = DrillConfig.create();
    try (final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        final Drillbit bit = new Drillbit(config, serviceSet)) {
        bit.run();
        final DrillbitContext context = bit.getContext();
        final MaterializedField intField = MaterializedField.create("int", Types.required(TypeProtos.MinorType.INT));
        final MaterializedField binField = MaterializedField.create("binary", Types.required(TypeProtos.MinorType.VARBINARY));
        try (final IntVector intVector = (IntVector) TypeHelper.getNewVector(intField, context.getAllocator());
            final VarBinaryVector binVector = (VarBinaryVector) TypeHelper.getNewVector(binField, context.getAllocator())) {
            AllocationHelper.allocate(intVector, 4, 4);
            AllocationHelper.allocate(binVector, 4, 5);
            vectorList.add(intVector);
            vectorList.add(binVector);
            intVector.getMutator().setSafe(0, 0);
            binVector.getMutator().setSafe(0, "ZERO".getBytes());
            intVector.getMutator().setSafe(1, 1);
            binVector.getMutator().setSafe(1, "ONE".getBytes());
            intVector.getMutator().setSafe(2, 2);
            binVector.getMutator().setSafe(2, "TWO".getBytes());
            intVector.getMutator().setSafe(3, 3);
            binVector.getMutator().setSafe(3, "THREE".getBytes());
            intVector.getMutator().setValueCount(4);
            binVector.getMutator().setValueCount(4);
            VectorContainer container = new VectorContainer();
            container.addCollection(vectorList);
            container.setRecordCount(4);
            @SuppressWarnings("resource") WritableBatch batch = WritableBatch.getBatchNoHVWrap(container.getRecordCount(), container, false);
            VectorAccessibleSerializable wrap = new VectorAccessibleSerializable(batch, context.getAllocator());
            final VectorAccessibleSerializable newWrap = new VectorAccessibleSerializable(context.getAllocator());
            try (final FileSystem fs = getLocalFileSystem()) {
                final File tempDir = Files.createTempDir();
                tempDir.deleteOnExit();
                final Path path = new Path(tempDir.getAbsolutePath(), "drillSerializable");
                try (final FSDataOutputStream out = fs.create(path)) {
                    wrap.writeToStream(out);
                }
                try (final FSDataInputStream in = fs.open(path)) {
                    newWrap.readFromStream(in);
                }
            }
            final VectorAccessible newContainer = newWrap.get();
            for (VectorWrapper<?> w : newContainer) {
                try (ValueVector vv = w.getValueVector()) {
                    int values = vv.getAccessor().getValueCount();
                    for (int i = 0; i < values; i++) {
                        final Object o = vv.getAccessor().getObject(i);
                        if (o instanceof byte[]) {
                            System.out.println(new String((byte[]) o));
                        } else {
                            System.out.println(o);
                        }
                    }
                }
            }
        }
    }
}

Example 8 with Drillbit

use of org.apache.drill.exec.server.Drillbit in project drill by axbaretto.

the class TestMultiInputAdd method testMultiInputAdd.

@Test
public void testMultiInputAdd() throws Throwable {
    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(DrillFileUtils.getResourceAsFile("/functions/multi_input_add_test.json"), Charsets.UTF_8));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        for (VectorWrapper<?> v : batchLoader) {
            ValueVector.Accessor accessor = v.getValueVector().getAccessor();
            assertTrue((accessor.getObject(0)).equals(10));
        }
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}

Example 9 with Drillbit

use of org.apache.drill.exec.server.Drillbit in project drill by axbaretto.

the class TestSimpleTopN method sortOneKeyAscending.

@Test
public void sortOneKeyAscending() throws Throwable {
    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<QueryDataBatch> results = client.runQuery(org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/topN/one_key_sort.json"), Charsets.UTF_8));
        int count = 0;
        for (QueryDataBatch b : results) {
            if (b.getHeader().getRowCount() != 0) {
                count += b.getHeader().getRowCount();
            }
        }
        assertEquals(100, count);
        long previousBigInt = Long.MIN_VALUE;
        int recordCount = 0;
        int batchCount = 0;
        for (QueryDataBatch b : results) {
            if (b.getHeader().getRowCount() == 0) {
                continue;
            }
            batchCount++;
            RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator());
            loader.load(b.getHeader().getDef(), b.getData());
            BigIntVector c1 = (BigIntVector) loader.getValueAccessorById(BigIntVector.class, loader.getValueVectorId(new SchemaPath("blue", ExpressionPosition.UNKNOWN)).getFieldIds()).getValueVector();
            BigIntVector.Accessor a1 = c1.getAccessor();
            for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
                recordCount++;
                assertTrue(previousBigInt <= a1.get(i));
                previousBigInt = a1.get(i);
            }
            loader.clear();
            b.release();
        }
        System.out.println(String.format("Sorted %,d records in %d batches.", recordCount, batchCount));
    }
}

Example 10 with Drillbit

use of org.apache.drill.exec.server.Drillbit in project drill by axbaretto.

the class TestBroadcastExchange method TestSingleBroadcastExchangeWithTwoScans.

@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(DrillFileUtils.getResourceAsFile("/sender/broadcast_exchange.json"), Charsets.UTF_8).replace("#{LEFT_FILE}", DrillFileUtils.getResourceAsFile("/join/merge_single_batch.left.json").toURI().toString()).replace("#{RIGHT_FILE}", DrillFileUtils.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);
    }
}