Examples with RowSetBuilder org.apache.drill.test.rowSet.RowSetBuilder used on opensource projects

Example 1 with RowSetBuilder

use of org.apache.drill.test.rowSet.RowSetBuilder in project drill by axbaretto.

the class TopNBatchTest method priorityQueueOrderingTest.

/**
 * Priority queue unit test.
 * @throws Exception
 */
@Test
public void priorityQueueOrderingTest() throws Exception {
    Properties properties = new Properties();
    DrillConfig drillConfig = DrillConfig.create(properties);
    FieldReference expr = FieldReference.getWithQuotedRef("colA");
    Order.Ordering ordering = new Order.Ordering(Order.Ordering.ORDER_DESC, expr, Order.Ordering.NULLS_FIRST);
    List<Order.Ordering> orderings = Lists.newArrayList(ordering);
    MaterializedField colA = MaterializedField.create("colA", Types.required(TypeProtos.MinorType.INT));
    MaterializedField colB = MaterializedField.create("colB", Types.required(TypeProtos.MinorType.INT));
    List<MaterializedField> cols = Lists.newArrayList(colA, colB);
    BatchSchema batchSchema = new BatchSchema(BatchSchema.SelectionVectorMode.NONE, cols);
    RowSet expectedRowSet;
    try (RootAllocator allocator = new RootAllocator(100_000_000)) {
        expectedRowSet = new RowSetBuilder(allocator, batchSchema).addRow(110, 10).addRow(109, 9).addRow(108, 8).addRow(107, 7).addRow(106, 6).addRow(105, 5).addRow(104, 4).addRow(103, 3).addRow(102, 2).addRow(101, 1).build();
        PriorityQueue queue;
        ExpandableHyperContainer hyperContainer;
        {
            VectorContainer container = new RowSetBuilder(allocator, batchSchema).build().container();
            hyperContainer = new ExpandableHyperContainer(container);
            queue = TopNBatch.createNewPriorityQueue(TopNBatch.createMainMappingSet(), TopNBatch.createLeftMappingSet(), TopNBatch.createRightMappingSet(), optionManager, new FunctionImplementationRegistry(drillConfig), new CodeCompiler(drillConfig, optionManager), orderings, hyperContainer, false, true, 10, allocator, batchSchema.getSelectionVectorMode());
        }
        List<RecordBatchData> testBatches = Lists.newArrayList();
        try {
            final Random random = new Random();
            final int bound = 100;
            final int numBatches = 11;
            final int numRecordsPerBatch = 100;
            for (int batchCounter = 0; batchCounter < numBatches; batchCounter++) {
                RowSetBuilder rowSetBuilder = new RowSetBuilder(allocator, batchSchema);
                rowSetBuilder.addRow((batchCounter + bound), batchCounter);
                for (int recordCounter = 0; recordCounter < numRecordsPerBatch; recordCounter++) {
                    rowSetBuilder.addRow(random.nextInt(bound), random.nextInt(bound));
                }
                VectorContainer vectorContainer = rowSetBuilder.build().container();
                queue.add(new RecordBatchData(vectorContainer, allocator));
            }
            queue.generate();
            VectorContainer resultContainer = queue.getHyperBatch();
            resultContainer.buildSchema(BatchSchema.SelectionVectorMode.NONE);
            RowSet.HyperRowSet actualHyperSet = HyperRowSetImpl.fromContainer(resultContainer, queue.getFinalSv4());
            new RowSetComparison(expectedRowSet).verify(actualHyperSet);
        } finally {
            if (expectedRowSet != null) {
                expectedRowSet.clear();
            }
            queue.cleanup();
            hyperContainer.clear();
            for (RecordBatchData testBatch : testBatches) {
                testBatch.clear();
            }
        }
    }
}

Example 2 with RowSetBuilder

use of org.apache.drill.test.rowSet.RowSetBuilder in project drill by axbaretto.

the class TestExternalSort method testNumericTypes.

/**
 * Test union type support in sort using numeric types: BIGINT and FLOAT8
 * Drill does not support union types fully. Sort was adapted to handle them.
 * This test simply verifies that the sort handles these types, even though
 * Drill does not.
 *
 * @param testLegacy
 *          true to test the old (pre-1.11) sort, false to test the new (1.11
 *          and later) sort
 * @throws Exception
 */
private void testNumericTypes(boolean testLegacy) throws Exception {
    final int record_count = 10000;
    final String tableDirName = "numericTypes";
    {
        final BatchSchema schema = new SchemaBuilder().add("a", Types.required(TypeProtos.MinorType.INT)).build();
        final RowSetBuilder rowSetBuilder = new RowSetBuilder(allocator, schema);
        for (int i = 0; i <= record_count; i += 2) {
            rowSetBuilder.addRow(i);
        }
        final RowSet rowSet = rowSetBuilder.build();
        final File tableFile = createTableFile(tableDirName, "a.json");
        new JsonFileBuilder(rowSet).build(tableFile);
        rowSet.clear();
    }
    {
        final BatchSchema schema = new SchemaBuilder().add("a", Types.required(TypeProtos.MinorType.FLOAT4)).build();
        final RowSetBuilder rowSetBuilder = new RowSetBuilder(allocator, schema);
        for (int i = 1; i <= record_count; i += 2) {
            rowSetBuilder.addRow((float) i);
        }
        final RowSet rowSet = rowSetBuilder.build();
        final File tableFile = createTableFile(tableDirName, "b.json");
        new JsonFileBuilder(rowSet).setCustomFormatter("a", "%.2f").build(tableFile);
        rowSet.clear();
    }
    TestBuilder builder = testBuilder().sqlQuery("select * from dfs.`%s` order by a desc", tableDirName).optionSettingQueriesForTestQuery(getOptions(testLegacy)).ordered().baselineColumns("a");
    for (int i = record_count; i >= 0; ) {
        builder.baselineValues((long) i--);
        if (i >= 0) {
            builder.baselineValues((double) i--);
        }
    }
    builder.go();
}

Example 3 with RowSetBuilder

use of org.apache.drill.test.rowSet.RowSetBuilder in project drill by axbaretto.

the class AbstractGenericCopierTest method testCopyRecords.

@Test
public void testCopyRecords() throws SchemaChangeException {
    try (RootAllocator allocator = new RootAllocator(10_000_000)) {
        final BatchSchema batchSchema = createTestSchema(BatchSchema.SelectionVectorMode.NONE);
        final RowSet srcRowSet = createSrcRowSet(allocator);
        final RowSet destRowSet = new RowSetBuilder(allocator, batchSchema).build();
        final VectorContainer destContainer = destRowSet.container();
        final Copier copier = createCopier();
        final RowSet expectedRowSet = createExpectedRowset(allocator);
        copier.setup(new RowSetBatch(srcRowSet), destContainer);
        copier.copyRecords(0, 3);
        try {
            new RowSetComparison(expectedRowSet).verify(destRowSet);
        } finally {
            srcRowSet.clear();
            if (srcRowSet instanceof RowSet.HyperRowSet) {
                ((RowSet.HyperRowSet) srcRowSet).getSv4().clear();
            }
            destRowSet.clear();
            expectedRowSet.clear();
        }
    }
}

Example 4 with RowSetBuilder

use of org.apache.drill.test.rowSet.RowSetBuilder in project drill by axbaretto.

the class TestSorter method testTwoRows.

// Paranoia: sort with two rows.
@Test
public void testTwoRows() throws Exception {
    BatchSchema schema = SortTestUtilities.nonNullSchema();
    SingleRowSet rowSet = new RowSetBuilder(fixture.allocator(), schema).addRow(1, "1").addRow(0, "0").withSv2().build();
    SingleRowSet expected = new RowSetBuilder(fixture.allocator(), schema).addRow(0, "0").addRow(1, "1").build();
    runSorterTest(rowSet, expected);
}

Example 5 with RowSetBuilder

use of org.apache.drill.test.rowSet.RowSetBuilder in project drill by axbaretto.

the class TestSorter method testSingleRow.

// Sanity test: single row
@Test
public void testSingleRow() throws Exception {
    BatchSchema schema = SortTestUtilities.nonNullSchema();
    SingleRowSet rowSet = new RowSetBuilder(fixture.allocator(), schema).addRow(0, "0").withSv2().build();
    SingleRowSet expected = new RowSetBuilder(fixture.allocator(), schema).addRow(0, "0").build();
    runSorterTest(rowSet, expected);
}