Examples with RecordBatchLoader org.apache.drill.exec.record.RecordBatchLoader used on opensource projects

Example 61 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestOrderedPartitionExchange method twoBitTwoExchangeRun.

/**
   * Starts two drillbits and runs a physical plan with a Mock scan, project, OrderedParititionExchange, Union Exchange,
   * and sort. The final sort is done first on the partition column, and verifies that the partitions are correct, in that
   * all rows in partition 0 should come in the sort order before any row in partition 1, etc. Also verifies that the standard
   * deviation of the size of the partitions is less than one tenth the mean size of the partitions, because we expect all
   * the partitions to be roughly equal in size.
   * @throws Exception
   */
@Test
public void twoBitTwoExchangeRun() 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<QueryDataBatch> results = client.runQuery(org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/sender/ordered_exchange.json"), Charsets.UTF_8));
        int count = 0;
        List<Integer> partitionRecordCounts = Lists.newArrayList();
        for (QueryDataBatch b : results) {
            if (b.getData() != null) {
                int rows = b.getHeader().getRowCount();
                count += rows;
                DrillConfig config = DrillConfig.create();
                RecordBatchLoader loader = new RecordBatchLoader(new BootStrapContext(config, ClassPathScanner.fromPrescan(config)).getAllocator());
                loader.load(b.getHeader().getDef(), b.getData());
                BigIntVector vv1 = (BigIntVector) loader.getValueAccessorById(BigIntVector.class, loader.getValueVectorId(new SchemaPath("col1", ExpressionPosition.UNKNOWN)).getFieldIds()).getValueVector();
                Float8Vector vv2 = (Float8Vector) loader.getValueAccessorById(Float8Vector.class, loader.getValueVectorId(new SchemaPath("col2", ExpressionPosition.UNKNOWN)).getFieldIds()).getValueVector();
                IntVector pVector = (IntVector) loader.getValueAccessorById(IntVector.class, loader.getValueVectorId(new SchemaPath("partition", ExpressionPosition.UNKNOWN)).getFieldIds()).getValueVector();
                long previous1 = Long.MIN_VALUE;
                double previous2 = Double.MIN_VALUE;
                int partPrevious = -1;
                long current1 = Long.MIN_VALUE;
                double current2 = Double.MIN_VALUE;
                int partCurrent = -1;
                int partitionRecordCount = 0;
                for (int i = 0; i < rows; i++) {
                    previous1 = current1;
                    previous2 = current2;
                    partPrevious = partCurrent;
                    current1 = vv1.getAccessor().get(i);
                    current2 = vv2.getAccessor().get(i);
                    partCurrent = pVector.getAccessor().get(i);
                    Assert.assertTrue(current1 >= previous1);
                    if (current1 == previous1) {
                        Assert.assertTrue(current2 <= previous2);
                    }
                    if (partCurrent == partPrevious || partPrevious == -1) {
                        partitionRecordCount++;
                    } else {
                        partitionRecordCounts.add(partitionRecordCount);
                        partitionRecordCount = 0;
                    }
                }
                partitionRecordCounts.add(partitionRecordCount);
                loader.clear();
            }
            b.release();
        }
        double[] values = new double[partitionRecordCounts.size()];
        int i = 0;
        for (Integer rc : partitionRecordCounts) {
            values[i++] = rc.doubleValue();
        }
        StandardDeviation stdDev = new StandardDeviation();
        Mean mean = new Mean();
        double std = stdDev.evaluate(values);
        double m = mean.evaluate(values);
        System.out.println("mean: " + m + " std dev: " + std);
        //Assert.assertTrue(std < 0.1 * m);
        assertEquals(31000, count);
    }
}

Example 62 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestDateTypes method testTimeStamp.

@Test
public void testTimeStamp() 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("/record/vector/test_timestamp.json"), Charsets.UTF_8).replace("#{TEST_FILE}", "/test_simple_date.json"));
        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();
            assertEquals(accessor.getObject(0).toString(), "1970-01-02 10:20:33.000");
            assertEquals(accessor.getObject(1).toString(), "2008-12-28 11:34:00.129");
            assertEquals(accessor.getObject(2).toString(), "2000-02-27 14:24:00.000");
        }
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}

Example 63 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestDateTypes method testInterval.

@Test
public void testInterval() 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("/record/vector/test_interval.json"), Charsets.UTF_8).replace("#{TEST_FILE}", "/test_simple_interval.json"));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        Iterator<VectorWrapper<?>> itr = batchLoader.iterator();
        ValueVector.Accessor accessor = itr.next().getValueVector().getAccessor();
        // Check the interval type
        assertEquals((accessor.getObject(0).toString()), ("2 years 2 months 1 day 1:20:35.0"));
        assertEquals((accessor.getObject(1).toString()), ("2 years 2 months 0 days 0:0:0.0"));
        assertEquals((accessor.getObject(2).toString()), ("0 years 0 months 0 days 1:20:35.0"));
        assertEquals((accessor.getObject(3).toString()), ("2 years 2 months 1 day 1:20:35.897"));
        assertEquals((accessor.getObject(4).toString()), ("0 years 0 months 0 days 0:0:35.4"));
        assertEquals((accessor.getObject(5).toString()), ("1 year 10 months 1 day 0:-39:-25.0"));
        accessor = itr.next().getValueVector().getAccessor();
        // Check the interval year type
        assertEquals((accessor.getObject(0).toString()), ("2 years 2 months "));
        assertEquals((accessor.getObject(1).toString()), ("2 years 2 months "));
        assertEquals((accessor.getObject(2).toString()), ("0 years 0 months "));
        assertEquals((accessor.getObject(3).toString()), ("2 years 2 months "));
        assertEquals((accessor.getObject(4).toString()), ("0 years 0 months "));
        assertEquals((accessor.getObject(5).toString()), ("1 year 10 months "));
        accessor = itr.next().getValueVector().getAccessor();
        // Check the interval day type
        assertEquals((accessor.getObject(0).toString()), ("1 day 1:20:35.0"));
        assertEquals((accessor.getObject(1).toString()), ("0 days 0:0:0.0"));
        assertEquals((accessor.getObject(2).toString()), ("0 days 1:20:35.0"));
        assertEquals((accessor.getObject(3).toString()), ("1 day 1:20:35.897"));
        assertEquals((accessor.getObject(4).toString()), ("0 days 0:0:35.4"));
        assertEquals((accessor.getObject(5).toString()), ("1 day 0:-39:-25.0"));
        batchLoader.clear();
        for (QueryDataBatch b : results) {
            b.release();
        }
    }
}

Example 64 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestLoad method testLoadValueVector.

@Test
public void testLoadValueVector() throws Exception {
    final BufferAllocator allocator = RootAllocatorFactory.newRoot(drillConfig);
    final ValueVector fixedV = new IntVector(MaterializedField.create("ints", Types.required(MinorType.INT)), allocator);
    final ValueVector varlenV = new VarCharVector(MaterializedField.create("chars", Types.required(MinorType.VARCHAR)), allocator);
    final ValueVector nullableVarlenV = new NullableVarCharVector(MaterializedField.create("chars", Types.optional(MinorType.VARCHAR)), allocator);
    final List<ValueVector> vectors = Lists.newArrayList(fixedV, varlenV, nullableVarlenV);
    for (final ValueVector v : vectors) {
        AllocationHelper.allocate(v, 100, 50);
        v.getMutator().generateTestData(100);
    }
    final WritableBatch writableBatch = WritableBatch.getBatchNoHV(100, vectors, false);
    final RecordBatchLoader batchLoader = new RecordBatchLoader(allocator);
    final ByteBuf[] byteBufs = writableBatch.getBuffers();
    int bytes = 0;
    for (int i = 0; i < byteBufs.length; i++) {
        bytes += byteBufs[i].writerIndex();
    }
    final DrillBuf byteBuf = allocator.buffer(bytes);
    int index = 0;
    for (int i = 0; i < byteBufs.length; i++) {
        byteBufs[i].readBytes(byteBuf, index, byteBufs[i].writerIndex());
        index += byteBufs[i].writerIndex();
    }
    byteBuf.writerIndex(bytes);
    batchLoader.load(writableBatch.getDef(), byteBuf);
    boolean firstColumn = true;
    int recordCount = 0;
    for (final VectorWrapper<?> v : batchLoader) {
        if (firstColumn) {
            firstColumn = false;
        } else {
            System.out.print("\t");
        }
        System.out.print(v.getField().getPath());
        System.out.print("[");
        System.out.print(v.getField().getType().getMinorType());
        System.out.print("]");
    }
    System.out.println();
    for (int r = 0; r < batchLoader.getRecordCount(); r++) {
        boolean first = true;
        recordCount++;
        for (final VectorWrapper<?> v : batchLoader) {
            if (first) {
                first = false;
            } else {
                System.out.print("\t");
            }
            final ValueVector.Accessor accessor = v.getValueVector().getAccessor();
            if (v.getField().getType().getMinorType() == TypeProtos.MinorType.VARCHAR) {
                final Object obj = accessor.getObject(r);
                if (obj != null) {
                    System.out.print(accessor.getObject(r));
                } else {
                    System.out.print("NULL");
                }
            } else {
                System.out.print(accessor.getObject(r));
            }
        }
        if (!first) {
            System.out.println();
        }
    }
    assertEquals(100, recordCount);
    batchLoader.clear();
    writableBatch.clear();
}

Example 65 with RecordBatchLoader

use of org.apache.drill.exec.record.RecordBatchLoader in project drill by apache.

the class TestExtractFunctions method testFrom.

private void testFrom(String fromType, String testDataFile, String columnName, long[][] expectedValues) 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("/functions/extractFrom.json"), Charsets.UTF_8).replace("#{TEST_TYPE}", fromType).replace("#{TEST_FILE}", testDataFile).replace("#{COLUMN_NAME}", columnName));
        RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());
        QueryDataBatch batch = results.get(0);
        assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));
        for (int i = 0; i < expectedValues.length; i++) {
            for (int j = 0; j < expectedValues[i].length; j++) {
                NullableBigIntVector vv = (NullableBigIntVector) batchLoader.getValueAccessorById(NullableBigIntVector.class, j).getValueVector();
                System.out.println("[" + i + "][" + j + "]: Expected: " + expectedValues[i][j] + ", Actual: " + vv.getAccessor().get(i));
                assertEquals(expectedValues[i][j], vv.getAccessor().get(i));
            }
        }
        for (QueryDataBatch b : results) {
            b.release();
        }
        batchLoader.clear();
    }
}