Examples with MockRecordBatch org.apache.drill.exec.physical.impl.MockRecordBatch used on opensource projects

Example 96 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch in project drill by apache.

the class TestLateralJoinCorrectness method test1RecordLeftBatchTo2RightRecordBatch.

/**
 * Test to show correct IterOutcome & output produced by LATERAL when one record in left batch produces 2 right
 * batches with OK and EMIT outcome respectively. Then output of LATERAL should be produced with OK outcome after the
 * join is done using both right batch with the left batch. It verifies the number of records in the output batch
 * based on left and right input batches.
 *
 * @throws Exception
 */
@Test
public void test1RecordLeftBatchTo2RightRecordBatch() throws Exception {
    // Get the left container with dummy data for Lateral Join
    leftContainer.add(nonEmptyLeftRowSet.container());
    // Get the left IterOutcomes for Lateral Join
    leftOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Create Left MockRecordBatch
    final CloseableRecordBatch leftMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, leftContainer, leftOutcomes, leftContainer.get(0).getSchema());
    // Get the right container with dummy data
    final RowSet.SingleRowSet nonEmptyRightRowSet2 = fixture.rowSetBuilder(rightSchema).addRow(1, 4, 41, "item41").addRow(1, 5, 51, "item51").build();
    rightContainer.add(emptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet.container());
    rightContainer.add(nonEmptyRightRowSet2.container());
    rightOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    rightOutcomes.add(RecordBatch.IterOutcome.OK);
    rightOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final CloseableRecordBatch rightMockBatch = new MockRecordBatch(fixture.getFragmentContext(), operatorContext, rightContainer, rightOutcomes, rightContainer.get(0).getSchema());
    final LateralJoinBatch ljBatch = new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), leftMockBatch, rightMockBatch);
    try {
        assertTrue(RecordBatch.IterOutcome.OK_NEW_SCHEMA == ljBatch.next());
        assertTrue(RecordBatch.IterOutcome.OK == ljBatch.next());
        assertTrue(ljBatch.getRecordCount() == (nonEmptyLeftRowSet.rowCount() * (nonEmptyRightRowSet.rowCount() + nonEmptyRightRowSet2.rowCount())));
        assertTrue(RecordBatch.IterOutcome.NONE == ljBatch.next());
    } catch (AssertionError | Exception error) {
        fail();
    } finally {
        // Close all the resources for this test case
        ljBatch.close();
        leftMockBatch.close();
        rightMockBatch.close();
        nonEmptyRightRowSet2.clear();
    }
}

Example 97 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch in project drill by apache.

the class TestHashAggEmitOutcome method testHashAggrEmit.

// private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestHashAggEmitOutcome.class);
/**
 *  A generic method to execute a Hash-Aggr emit test, based on the given parameters.
 *  Can take at most two generic input batches, and verify against at most two non-empty output
 *  batches (and unlimited number of input/output empty batches may be used)
 *
 *      This interface is a little ugly, because Java does not support simple initializations
 *      other than for arrays (e.g., no "structs" like in c++)
 *
 *   Input batch 1 is build in (see BaseTestOpBatchEmitOutcome.java)
 * @param inp2_1 - Input batch 2, first col (use null if not needed)
 * @param inp2_2 - input batch 2, second col
 * @param inp2_3 - Input batch 2, third col
 * @param inp3_1 - Input batch 3, first col (use null if not needed)
 * @param inp3_2 - input batch 3, second col
 * @param inp3_3 - input batch 3, third col
 * @param exp1_1 - First expected batch, col 1
 * @param exp1_2 - First expected batch, col 2
 * @param exp2_1 - Second expected batch, col 1
 * @param exp2_2 - Second expected batch, col 2
 * @param inpRowSet - Which input batches to use (the empty, i.e. 0, can be used multiple times)
 * @param inpOutcomes - Which input IterOutcomes to mark each input batch
 * @param outputRowCounts - expected number of rows, in each output batch
 * @param outputOutcomes - the expected output outcomes
 */
private // first input batch
void testHashAggrEmit(// first input batch
int[] inp2_1, // first input batch
int[] inp2_2, // first input batch
String[] inp2_3, // second input batch
int[] inp3_1, // second input batch
int[] inp3_2, // second input batch
String[] inp3_3, // first expected
String[] exp1_1, // first expected
int[] exp1_2, // second expected
String[] exp2_1, // second expected
int[] exp2_2, // input batches + outcomes
int[] inpRowSet, // input batches + outcomes
RecordBatch.IterOutcome[] inpOutcomes, // output row counts per each out batch
List<Integer> outputRowCounts, // output outcomes
List<RecordBatch.IterOutcome> outputOutcomes) {
    // First input batch
    RowSetBuilder builder2 = operatorFixture.rowSetBuilder(inputSchema);
    if (inp2_1 != null) {
        for (int i = 0; i < inp2_1.length; i++) {
            builder2 = builder2.addRow(inp2_1[i], inp2_2[i], inp2_3[i]);
        }
    }
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = builder2.build();
    // Second input batch
    RowSetBuilder builder3 = operatorFixture.rowSetBuilder(inputSchema);
    if (inp3_1 != null) {
        for (int i = 0; i < inp3_1.length; i++) {
            builder3 = builder3.addRow(inp3_1[i], inp3_2[i], inp3_3[i]);
        }
    }
    final RowSet.SingleRowSet nonEmptyInputRowSet3 = builder3.build();
    final TupleMetadata resultSchema = new SchemaBuilder().add("name", TypeProtos.MinorType.VARCHAR).addNullable("total_sum", TypeProtos.MinorType.BIGINT).buildSchema();
    // First expected batch
    RowSetBuilder expectedBuilder1 = operatorFixture.rowSetBuilder(resultSchema);
    if (exp1_1 != null) {
        for (int i = 0; i < exp1_1.length; i++) {
            expectedBuilder1 = expectedBuilder1.addRow(exp1_1[i], (long) exp1_2[i]);
        }
    }
    final RowSet.SingleRowSet expectedRowSet1 = expectedBuilder1.build();
    // Second expected batch
    RowSetBuilder expectedBuilder2 = operatorFixture.rowSetBuilder(resultSchema);
    if (exp2_1 != null) {
        for (int i = 0; i < exp2_1.length; i++) {
            expectedBuilder2 = expectedBuilder2.addRow(exp2_1[i], (long) exp2_2[i]);
        }
    }
    final RowSet.SingleRowSet expectedRowSet2 = expectedBuilder2.build();
    // Add the input batches, in the order/type given
    for (int inp : inpRowSet) {
        switch(inp) {
            case 0:
                inputContainer.add(emptyInputRowSet.container());
                break;
            case 1:
                inputContainer.add(nonEmptyInputRowSet.container());
                break;
            case 2:
                inputContainer.add(nonEmptyInputRowSet2.container());
                break;
            case 3:
                inputContainer.add(nonEmptyInputRowSet3.container());
                break;
            default:
                fail();
        }
    }
    // build the outcomes
    inputOutcomes.addAll(Arrays.asList(inpOutcomes));
    // 
    // Build the Hash Agg Batch operator
    // 
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final HashAggregate hashAggrConfig = new HashAggregate(null, AggPrelBase.OperatorPhase.PHASE_1of1, parseExprs("name_left", "name"), parseExprs("sum(id_left+cost_left)", "total_sum"), 1.0f);
    final HashAggBatch haBatch = new HashAggBatch(hashAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    // 
    // Iterate thru the next batches, and verify expected outcomes
    // 
    assertEquals(outputRowCounts.size(), outputOutcomes.size());
    boolean firstOne = true;
    for (int ind = 0; ind < outputOutcomes.size(); ind++) {
        RecordBatch.IterOutcome expOut = outputOutcomes.get(ind);
        assertSame(expOut, haBatch.next());
        if (expOut == NONE) {
            break;
        }
        // done
        RowSet actualRowSet = DirectRowSet.fromContainer(haBatch.getContainer());
        int expectedSize = outputRowCounts.get(ind);
        // System.out.println(expectedSize);
        if (0 == expectedSize) {
            assertEquals(expectedSize, haBatch.getRecordCount());
        } else if (firstOne) {
            firstOne = false;
            new RowSetComparison(expectedRowSet1).verify(actualRowSet);
        } else {
            new RowSetComparison(expectedRowSet2).verify(actualRowSet);
        }
    }
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    nonEmptyInputRowSet3.clear();
    expectedRowSet2.clear();
    expectedRowSet1.clear();
}

Example 98 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t12_testStreamingAggrNonEmptyBatchEmitOutcome.

/**
 *   Repeats t2_testStreamingAggrNonEmptyBatchEmitOutcome with no group by
 */
@Test
public void t12_testStreamingAggrNonEmptyBatchEmitOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(13, 130, "item13").addRow(13, 130, "item13").addRow(2, 20, "item2").addRow(2, 20, "item2").addRow(4, 40, "item4").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchemaNoGroupBy).addRow((long) 385).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<NamedExpression>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    // Data before EMIT is returned with an OK_NEW_SCHEMA.
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    // EMIT comes with an empty batch
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    // Release memory for row sets
    nonEmptyInputRowSet2.clear();
    expectedRowSet.clear();
}

Example 99 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t7_testStreamingAggrMultipleEMITOutcome.

/**
 * Normal case
 */
@Test
public void t7_testStreamingAggrMultipleEMITOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item2").addRow(3, 30, "item3").build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    inputOutcomes.add(RecordBatch.IterOutcome.EMIT);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, parseExprs("name_left", "name"), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(2, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.EMIT);
    assertEquals(0, strAggBatch.getRecordCount());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    nonEmptyInputRowSet2.clear();
}

Example 100 with MockRecordBatch

use of org.apache.drill.exec.physical.impl.MockRecordBatch in project drill by apache.

the class TestStreamingAggEmitOutcome method t19_testStreamingAgr_WithEmptyNonEmptyBatchesAndOKOutcome.

/**
 *   Repeats t9_testStreamingAgr_WithEmptyNonEmptyBatchesAndOKOutcome with no group by
 */
@Test
public void t19_testStreamingAgr_WithEmptyNonEmptyBatchesAndOKOutcome() {
    final RowSet.SingleRowSet nonEmptyInputRowSet2 = operatorFixture.rowSetBuilder(inputSchema).addRow(2, 20, "item1").addRow(13, 130, "item13").addRow(13, 130, "item13").addRow(13, 130, "item13").addRow(130, 1300, "item130").addRow(0, 0, "item130").build();
    final RowSet.SingleRowSet nonEmptyInputRowSet3 = operatorFixture.rowSetBuilder(inputSchema).addRow(23, 230, "item23").addRow(3, 33, "item3").addRow(7, 70, "item7").addRow(17, 170, "item7").build();
    final RowSet.SingleRowSet expectedRowSet = operatorFixture.rowSetBuilder(resultSchemaNoGroupBy).addRow((long) 2445).build();
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet2.container());
    inputContainer.add(emptyInputRowSet.container());
    inputContainer.add(nonEmptyInputRowSet3.container());
    inputContainer.add(emptyInputRowSet.container());
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    inputOutcomes.add(RecordBatch.IterOutcome.OK);
    final MockRecordBatch mockInputBatch = new MockRecordBatch(operatorFixture.getFragmentContext(), opContext, inputContainer, inputOutcomes, emptyInputRowSet.container().getSchema());
    final StreamingAggregate streamAggrConfig = new StreamingAggregate(null, new ArrayList<>(), parseExprs("sum(id_left+cost_left)", "total_sum"));
    final StreamingAggBatch strAggBatch = new StreamingAggBatch(streamAggrConfig, mockInputBatch, operatorFixture.getFragmentContext());
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.OK_NEW_SCHEMA);
    assertEquals(1, strAggBatch.getRecordCount());
    RowSet actualRowSet = DirectRowSet.fromContainer(strAggBatch.getContainer());
    new RowSetComparison(expectedRowSet).verify(actualRowSet);
    assertTrue(strAggBatch.next() == RecordBatch.IterOutcome.NONE);
    nonEmptyInputRowSet2.clear();
    nonEmptyInputRowSet3.clear();
    expectedRowSet.clear();
}