Examples with VariableReferenceExpression com.facebook.presto.spi.relation.VariableReferenceExpression used on opensource projects

Example 6 with VariableReferenceExpression

use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.

the class TestPinotQueryGenerator method testAggWithArrayFunctionsInGroupBy.

@Test
public void testAggWithArrayFunctionsInGroupBy() {
    LinkedHashMap<String, String> aggProjection = new LinkedHashMap<>();
    aggProjection.put("array_max_0", "array_max(scores)");
    PlanNode justMaxScores = buildPlan(planBuilder -> project(planBuilder, tableScan(planBuilder, pinotTable, regionId, secondsSinceEpoch, city, fare, scores), aggProjection, defaultSessionHolder));
    testPinotQuery(planBuilder -> planBuilder.aggregation(aggBuilder -> aggBuilder.source(justMaxScores).singleGroupingSet(new VariableReferenceExpression(Optional.empty(), "array_max_0", DOUBLE)).addAggregation(planBuilder.variable("agg"), getRowExpression("count(*)", defaultSessionHolder))), String.format("SELECT %s FROM realtimeOnly GROUP BY arrayMax(scores) %s 10000", getExpectedAggOutput("count(*)", "arrayMax(scores)"), getGroupByLimitKey()));
    aggProjection.put("city", "city");
    PlanNode newScanWithCity = buildPlan(planBuilder -> project(planBuilder, tableScan(planBuilder, pinotTable, regionId, secondsSinceEpoch, city, fare, scores), aggProjection, defaultSessionHolder));
    testPinotQuery(planBuilder -> planBuilder.aggregation(aggBuilder -> aggBuilder.source(newScanWithCity).singleGroupingSet(new VariableReferenceExpression(Optional.empty(), "array_max_0", DOUBLE), variable("city")).addAggregation(planBuilder.variable("agg"), getRowExpression("count(*)", defaultSessionHolder))), String.format("SELECT %s FROM realtimeOnly GROUP BY arrayMax(scores), city %s 10000", getExpectedAggOutput("count(*)", "arrayMax(scores), city"), getGroupByLimitKey()));
}

Example 7 with VariableReferenceExpression

use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.

the class TestPinotQueryGenerator method testDistinctCountPushdownWithVariableSuffix.

@Test
public void testDistinctCountPushdownWithVariableSuffix() {
    Map<VariableReferenceExpression, PinotColumnHandle> columnHandleMap = ImmutableMap.of(new VariableReferenceExpression(Optional.empty(), "regionid_33", regionId.getDataType()), regionId);
    PlanNode justScan = buildPlan(planBuilder -> tableScan(planBuilder, pinotTable, columnHandleMap));
    PlanNode markDistinct = buildPlan(planBuilder -> markDistinct(planBuilder, variable("regionid$distinct_62"), ImmutableList.of(variable("regionid")), justScan));
    PlanNode aggregate = buildPlan(planBuilder -> planBuilder.aggregation(aggBuilder -> aggBuilder.source(markDistinct).addAggregation(planBuilder.variable("count(regionid_33)"), getRowExpression("count(regionid_33)", defaultSessionHolder), Optional.empty(), Optional.empty(), false, Optional.of(variable("regionid$distinct_62"))).globalGrouping()));
    testPinotQuery(new PinotConfig().setAllowMultipleAggregations(true), planBuilder -> planBuilder.limit(10, aggregate), "SELECT DISTINCTCOUNT(regionId) FROM realtimeOnly");
}

Example 8 with VariableReferenceExpression

use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.

the class DruidPushdownUtils method computeAggregationNodes.

public static List<DruidAggregationColumnNode> computeAggregationNodes(AggregationNode aggregationNode) {
    int groupByKeyIndex = 0;
    ImmutableList.Builder<DruidAggregationColumnNode> nodeBuilder = ImmutableList.builder();
    for (VariableReferenceExpression outputColumn : aggregationNode.getOutputVariables()) {
        AggregationNode.Aggregation aggregation = aggregationNode.getAggregations().get(outputColumn);
        if (aggregation != null) {
            if (aggregation.getFilter().isPresent() || aggregation.isDistinct() || aggregation.getOrderBy().isPresent()) {
                throw new PrestoException(DRUID_PUSHDOWN_UNSUPPORTED_EXPRESSION, "Unsupported aggregation node " + aggregationNode);
            }
            if (aggregation.getMask().isPresent()) {
                // E.g. `SELECT count(distinct COL_A), sum(COL_B) FROM myTable` to Druid as `SELECT distinctCount(COL_A), sum(COL_B) FROM myTable`
                if (aggregation.getCall().getDisplayName().equalsIgnoreCase(COUNT_FUNCTION_NAME) && aggregation.getMask().get().getName().equalsIgnoreCase(aggregation.getArguments().get(0) + DISTINCT_MASK)) {
                    nodeBuilder.add(new AggregationFunctionColumnNode(outputColumn, new CallExpression(aggregation.getCall().getSourceLocation(), DRUID_COUNT_DISTINCT_FUNCTION_NAME, aggregation.getCall().getFunctionHandle(), aggregation.getCall().getType(), aggregation.getCall().getArguments())));
                    continue;
                }
                // Druid doesn't support push down aggregation functions other than count on top of distinct function.
                throw new PrestoException(DRUID_PUSHDOWN_UNSUPPORTED_EXPRESSION, "Unsupported aggregation node with mask " + aggregationNode);
            }
            if (handlePushDownSingleDistinctCount(nodeBuilder, aggregationNode, outputColumn, aggregation)) {
                continue;
            }
            nodeBuilder.add(new AggregationFunctionColumnNode(outputColumn, aggregation.getCall()));
        } else {
            // group by output
            VariableReferenceExpression inputColumn = aggregationNode.getGroupingKeys().get(groupByKeyIndex);
            nodeBuilder.add(new GroupByColumnNode(inputColumn, outputColumn));
            groupByKeyIndex++;
        }
    }
    return nodeBuilder.build();
}

Example 9 with VariableReferenceExpression

use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.

the class DruidQueryGeneratorContext method withAggregation.

public DruidQueryGeneratorContext withAggregation(Map<VariableReferenceExpression, Selection> newSelections, Map<VariableReferenceExpression, Selection> newGroupByColumns, int newAggregations, Set<VariableReferenceExpression> newHiddenColumnSet) {
    AtomicBoolean pushDownDistinctCount = new AtomicBoolean(false);
    newSelections.values().forEach(selection -> {
        if (selection.getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) {
            pushDownDistinctCount.set(true);
        }
    });
    Map<VariableReferenceExpression, Selection> targetSelections = newSelections;
    if (pushDownDistinctCount.get()) {
        // Push down count distinct query to Druid, clean up hidden column set by the non-aggregation groupBy Plan.
        newHiddenColumnSet = ImmutableSet.of();
        ImmutableMap.Builder<VariableReferenceExpression, Selection> builder = ImmutableMap.builder();
        for (Map.Entry<VariableReferenceExpression, Selection> entry : newSelections.entrySet()) {
            if (entry.getValue().getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) {
                String definition = entry.getValue().getDefinition();
                int start = definition.indexOf("(");
                int end = definition.indexOf(")");
                String countDistinctClause = "count ( distinct " + escapeSqlIdentifier(definition.substring(start + 1, end)) + ")";
                Selection countDistinctSelection = new Selection(countDistinctClause, entry.getValue().getOrigin());
                builder.put(entry.getKey(), countDistinctSelection);
            } else {
                builder.put(entry.getKey(), entry.getValue());
            }
        }
        targetSelections = builder.build();
    } else {
        checkState(!hasAggregation(), "Druid doesn't support aggregation on top of the aggregated data");
    }
    checkState(!hasLimit(), "Druid doesn't support aggregation on top of the limit");
    checkState(newAggregations > 0, "Invalid number of aggregations");
    return new DruidQueryGeneratorContext(targetSelections, from, filter, limit, newAggregations, newGroupByColumns, variablesInAggregation, newHiddenColumnSet, tableScanNodeId);
}

Example 10 with VariableReferenceExpression

use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.

the class SubfieldExtractor method toRowExpression.

private RowExpression toRowExpression(Subfield subfield, List<Type> types) {
    List<Subfield.PathElement> path = subfield.getPath();
    if (path.isEmpty()) {
        return new VariableReferenceExpression(Optional.empty(), subfield.getRootName(), types.get(0));
    }
    RowExpression base = toRowExpression(new Subfield(subfield.getRootName(), path.subList(0, path.size() - 1)), types.subList(0, types.size() - 1));
    Type baseType = types.get(types.size() - 2);
    Subfield.PathElement pathElement = path.get(path.size() - 1);
    if (pathElement instanceof Subfield.LongSubscript) {
        Type indexType = baseType instanceof MapType ? ((MapType) baseType).getKeyType() : BIGINT;
        FunctionHandle functionHandle = functionResolution.subscriptFunction(baseType, indexType);
        ConstantExpression index = new ConstantExpression(base.getSourceLocation(), ((Subfield.LongSubscript) pathElement).getIndex(), indexType);
        return new CallExpression(base.getSourceLocation(), SUBSCRIPT.name(), functionHandle, types.get(types.size() - 1), ImmutableList.of(base, index));
    }
    if (pathElement instanceof Subfield.StringSubscript) {
        Type indexType = ((MapType) baseType).getKeyType();
        FunctionHandle functionHandle = functionResolution.subscriptFunction(baseType, indexType);
        ConstantExpression index = new ConstantExpression(base.getSourceLocation(), Slices.utf8Slice(((Subfield.StringSubscript) pathElement).getIndex()), indexType);
        return new CallExpression(base.getSourceLocation(), SUBSCRIPT.name(), functionHandle, types.get(types.size() - 1), ImmutableList.of(base, index));
    }
    if (pathElement instanceof Subfield.NestedField) {
        Subfield.NestedField nestedField = (Subfield.NestedField) pathElement;
        return new SpecialFormExpression(base.getSourceLocation(), DEREFERENCE, types.get(types.size() - 1), base, new ConstantExpression(base.getSourceLocation(), getFieldIndex((RowType) baseType, nestedField.getName()), INTEGER));
    }
    verify(false, "Unexpected path element: " + pathElement);
    return null;
}