Examples with TransformExpressionTree com.linkedin.pinot.common.request.transform.TransformExpressionTree used on opensource projects

Example 1 with TransformExpressionTree

use of com.linkedin.pinot.common.request.transform.TransformExpressionTree 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);
    }
}

Example 2 with TransformExpressionTree

use of com.linkedin.pinot.common.request.transform.TransformExpressionTree in project pinot by linkedin.

the class TransformExpressionTreeTest method test.

/**
   * This test validates an expression tree built by {@link TransformExpressionTree#buildTree(AstNode)}
   */
@Test
public void test() {
    TransformFunctionFactory.init(new String[] { TransformFunctionFactoryTest.foo.class.getName(), TransformFunctionFactoryTest.bar.class.getName() });
    Pql2Compiler compiler = new Pql2Compiler();
    String expression = "foo(bar('a', foo(b, 'c', d)), e)";
    TransformExpressionTree expressionTree = compiler.compileToExpressionTree(expression);
    TransformFunction rootTransform = TransformFunctionFactory.get(expressionTree.getTransformName());
    Assert.assertEquals(rootTransform.getName(), "foo");
    List<TransformExpressionTree> firstChildren = expressionTree.getChildren();
    Assert.assertEquals(firstChildren.size(), 2);
    TransformExpressionTree firstChild = firstChildren.get(0);
    Assert.assertEquals(firstChild.getTransformName(), "bar");
    Assert.assertEquals(firstChildren.get(1).toString(), "e");
    List<TransformExpressionTree> secondChildren = firstChild.getChildren();
    Assert.assertEquals(secondChildren.size(), 2);
    Assert.assertEquals(secondChildren.get(0).toString(), "a");
    Assert.assertEquals(secondChildren.get(1).getTransformName(), "foo");
    List<TransformExpressionTree> thirdChildren = secondChildren.get(1).getChildren();
    Assert.assertEquals(thirdChildren.get(0).toString(), "b");
    Assert.assertEquals(thirdChildren.get(1).toString(), "c");
    Assert.assertEquals(thirdChildren.get(2).toString(), "d");
}

Example 3 with TransformExpressionTree

use of com.linkedin.pinot.common.request.transform.TransformExpressionTree in project pinot by linkedin.

the class DefaultExpressionEvaluator method evaluateExpression.

/**
   * Helper (recursive) method that walks the expression tree bottom up evaluating
   * transforms at each level.
   *
   * @param projectionBlock Projection block for which to evaluate the expression for
   * @param expressionTree Expression tree to evaluate
   * @return Result of the expression transform
   */
private BlockValSet evaluateExpression(ProjectionBlock projectionBlock, TransformExpressionTree expressionTree) {
    TransformFunction function = getTransformFunction(expressionTree.getTransformName());
    int numDocs = projectionBlock.getNumDocs();
    String expressionString = expressionTree.toString();
    TransformExpressionTree.ExpressionType expressionType = expressionTree.getExpressionType();
    switch(expressionType) {
        case FUNCTION:
            List<TransformExpressionTree> children = expressionTree.getChildren();
            int numChildren = children.size();
            BlockValSet[] transformArgs = new BlockValSet[numChildren];
            for (int i = 0; i < numChildren; i++) {
                transformArgs[i] = evaluateExpression(projectionBlock, children.get(i));
            }
            return new TransformBlockValSet(function, numDocs, transformArgs);
        case IDENTIFIER:
            return projectionBlock.getBlockValueSet(expressionString);
        case LITERAL:
            return new ConstantBlockValSet(expressionString, numDocs);
        default:
            throw new IllegalArgumentException("Illegal expression type in expression evaluator: " + expressionType);
    }
}

Example 4 with TransformExpressionTree

use of com.linkedin.pinot.common.request.transform.TransformExpressionTree in project pinot by linkedin.

the class TransformExpressionOperatorTest method evaluateExpression.

/**
   * Helper method to evaluate one expression using the TransformOperator.
   * @param expression Expression to evaluate
   * @return Result of evaluation
   */
private double[] evaluateExpression(String expression) {
    Operator filterOperator = new MatchEntireSegmentOperator(_indexSegment.getSegmentMetadata().getTotalDocs());
    final BReusableFilteredDocIdSetOperator docIdSetOperator = new BReusableFilteredDocIdSetOperator(filterOperator, _indexSegment.getSegmentMetadata().getTotalDocs(), NUM_ROWS);
    final Map<String, BaseOperator> dataSourceMap = buildDataSourceMap(_indexSegment.getSegmentMetadata().getSchema());
    final MProjectionOperator projectionOperator = new MProjectionOperator(dataSourceMap, docIdSetOperator);
    Pql2Compiler compiler = new Pql2Compiler();
    List<TransformExpressionTree> expressionTrees = new ArrayList<>(1);
    expressionTrees.add(compiler.compileToExpressionTree(expression));
    TransformExpressionOperator transformOperator = new TransformExpressionOperator(projectionOperator, expressionTrees);
    transformOperator.open();
    TransformBlock transformBlock = (TransformBlock) transformOperator.getNextBlock();
    BlockValSet blockValueSet = transformBlock.getBlockValueSet(expression);
    double[] actual = blockValueSet.getDoubleValuesSV();
    transformOperator.close();
    return actual;
}