Examples with GroupKeyGenerator com.linkedin.pinot.core.query.aggregation.groupby.GroupKeyGenerator used on opensource projects

Example 1 with GroupKeyGenerator

use of com.linkedin.pinot.core.query.aggregation.groupby.GroupKeyGenerator in project pinot by linkedin.

the class NoDictionaryGroupKeyGeneratorTest method testGroupKeyGenerator.

private void testGroupKeyGenerator(String[] groupByColumns, FieldSpec.DataType[] dataTypes) throws Exception {
    // Build the projection operator.
    MatchEntireSegmentOperator matchEntireSegmentOperator = new MatchEntireSegmentOperator(NUM_ROWS);
    BReusableFilteredDocIdSetOperator docIdSetOperator = new BReusableFilteredDocIdSetOperator(matchEntireSegmentOperator, NUM_ROWS, 10000);
    MProjectionOperator projectionOperator = new MProjectionOperator(_dataSourceMap, docIdSetOperator);
    TransformExpressionOperator transformOperator = new TransformExpressionOperator(projectionOperator, new ArrayList<TransformExpressionTree>());
    // Iterator over all projection blocks and generate group keys.
    TransformBlock transformBlock;
    int[] docIdToGroupKeys = new int[DocIdSetPlanNode.MAX_DOC_PER_CALL];
    GroupKeyGenerator groupKeyGenerator = null;
    while ((transformBlock = (TransformBlock) transformOperator.nextBlock()) != null) {
        if (groupKeyGenerator == null) {
            // Build the group key generator.
            groupKeyGenerator = (groupByColumns.length == 1) ? new NoDictionarySingleColumnGroupKeyGenerator(groupByColumns[0], dataTypes[0]) : new NoDictionaryMultiColumnGroupKeyGenerator(transformBlock, groupByColumns);
        }
        groupKeyGenerator.generateKeysForBlock(transformBlock, docIdToGroupKeys);
    }
    // Assert total number of group keys is as expected
    Assert.assertTrue(groupKeyGenerator != null);
    Set<String> expectedGroupKeys = getExpectedGroupKeys(_recordReader, groupByColumns);
    Assert.assertEquals(groupKeyGenerator.getCurrentGroupKeyUpperBound(), expectedGroupKeys.size(), "Number of group keys mis-match.");
    // Assert all group key values are as expected
    Iterator<GroupKeyGenerator.GroupKey> uniqueGroupKeys = groupKeyGenerator.getUniqueGroupKeys();
    while (uniqueGroupKeys.hasNext()) {
        GroupKeyGenerator.GroupKey groupKey = uniqueGroupKeys.next();
        String actual = groupKey.getStringKey();
        Assert.assertTrue(expectedGroupKeys.contains(actual), "Unexpected group key: " + actual);
    }
}