Examples with UnnestNode com.facebook.presto.sql.planner.plan.UnnestNode used on opensource projects

Example 1 with UnnestNode

use of com.facebook.presto.sql.planner.plan.UnnestNode in project presto by prestodb.

the class RelationPlanner method planCrossJoinUnnest.

private RelationPlan planCrossJoinUnnest(RelationPlan leftPlan, Join joinNode, Unnest node) {
    RelationType unnestOutputDescriptor = analysis.getOutputDescriptor(node);
    // Create variables for the result of unnesting
    ImmutableList.Builder<VariableReferenceExpression> unnestedVariablesBuilder = ImmutableList.builder();
    for (Field field : unnestOutputDescriptor.getVisibleFields()) {
        VariableReferenceExpression variable = variableAllocator.newVariable(field);
        unnestedVariablesBuilder.add(variable);
    }
    ImmutableList<VariableReferenceExpression> unnestedVariables = unnestedVariablesBuilder.build();
    // Add a projection for all the unnest arguments
    PlanBuilder planBuilder = initializePlanBuilder(leftPlan);
    planBuilder = planBuilder.appendProjections(node.getExpressions(), variableAllocator, idAllocator);
    TranslationMap translations = planBuilder.getTranslations();
    ProjectNode projectNode = (ProjectNode) planBuilder.getRoot();
    ImmutableMap.Builder<VariableReferenceExpression, List<VariableReferenceExpression>> unnestVariables = ImmutableMap.builder();
    UnmodifiableIterator<VariableReferenceExpression> unnestedVariablesIterator = unnestedVariables.iterator();
    for (Expression expression : node.getExpressions()) {
        Type type = analysis.getType(expression);
        VariableReferenceExpression inputVariable = new VariableReferenceExpression(getSourceLocation(expression), translations.get(expression).getName(), type);
        if (type instanceof ArrayType) {
            Type elementType = ((ArrayType) type).getElementType();
            if (!SystemSessionProperties.isLegacyUnnest(session) && elementType instanceof RowType) {
                ImmutableList.Builder<VariableReferenceExpression> unnestVariableBuilder = ImmutableList.builder();
                for (int i = 0; i < ((RowType) elementType).getFields().size(); i++) {
                    unnestVariableBuilder.add(unnestedVariablesIterator.next());
                }
                unnestVariables.put(inputVariable, unnestVariableBuilder.build());
            } else {
                unnestVariables.put(inputVariable, ImmutableList.of(unnestedVariablesIterator.next()));
            }
        } else if (type instanceof MapType) {
            unnestVariables.put(inputVariable, ImmutableList.of(unnestedVariablesIterator.next(), unnestedVariablesIterator.next()));
        } else {
            throw new IllegalArgumentException("Unsupported type for UNNEST: " + type);
        }
    }
    Optional<VariableReferenceExpression> ordinalityVariable = node.isWithOrdinality() ? Optional.of(unnestedVariablesIterator.next()) : Optional.empty();
    checkState(!unnestedVariablesIterator.hasNext(), "Not all output variables were matched with input variables");
    UnnestNode unnestNode = new UnnestNode(getSourceLocation(node), idAllocator.getNextId(), projectNode, leftPlan.getFieldMappings(), unnestVariables.build(), ordinalityVariable);
    return new RelationPlan(unnestNode, analysis.getScope(joinNode), unnestNode.getOutputVariables());
}

Example 2 with UnnestNode

use of com.facebook.presto.sql.planner.plan.UnnestNode in project presto by prestodb.

the class ExtractSpatialJoins method addPartitioningNodes.

private static PlanNode addPartitioningNodes(Context context, FunctionAndTypeManager functionAndTypeManager, PlanNode node, VariableReferenceExpression partitionVariable, KdbTree kdbTree, RowExpression geometry, Optional<RowExpression> radius) {
    Assignments.Builder projections = Assignments.builder();
    for (VariableReferenceExpression outputVariable : node.getOutputVariables()) {
        projections.put(outputVariable, outputVariable);
    }
    FunctionHandle castFunctionHandle = functionAndTypeManager.lookupCast(CAST, VARCHAR.getTypeSignature(), KDB_TREE.getTypeSignature());
    ImmutableList.Builder partitioningArgumentsBuilder = ImmutableList.builder().add(new CallExpression(partitionVariable.getSourceLocation(), CAST.name(), castFunctionHandle, KDB_TREE, ImmutableList.of(Expressions.constant(utf8Slice(KdbTreeUtils.toJson(kdbTree)), VARCHAR)))).add(geometry);
    radius.map(partitioningArgumentsBuilder::add);
    List<RowExpression> partitioningArguments = partitioningArgumentsBuilder.build();
    String spatialPartitionsFunctionName = "spatial_partitions";
    FunctionHandle functionHandle = functionAndTypeManager.lookupFunction(spatialPartitionsFunctionName, fromTypes(partitioningArguments.stream().map(RowExpression::getType).collect(toImmutableList())));
    CallExpression partitioningFunction = new CallExpression(partitionVariable.getSourceLocation(), spatialPartitionsFunctionName, functionHandle, new ArrayType(INTEGER), partitioningArguments);
    VariableReferenceExpression partitionsVariable = context.getVariableAllocator().newVariable(partitioningFunction);
    projections.put(partitionsVariable, partitioningFunction);
    return new UnnestNode(node.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(node.getSourceLocation(), context.getIdAllocator().getNextId(), node, projections.build(), LOCAL), node.getOutputVariables(), ImmutableMap.of(partitionsVariable, ImmutableList.of(partitionVariable)), Optional.empty());
}

Example 3 with UnnestNode

use of com.facebook.presto.sql.planner.plan.UnnestNode in project presto by prestodb.

the class CanonicalPlanGenerator method visitUnnest.

@Override
public Optional<PlanNode> visitUnnest(UnnestNode node, Map<VariableReferenceExpression, VariableReferenceExpression> context) {
    Optional<PlanNode> source = node.getSource().accept(this, context);
    if (!source.isPresent()) {
        return Optional.empty();
    }
    // Generate canonical unnestVariables.
    ImmutableMap.Builder<VariableReferenceExpression, List<VariableReferenceExpression>> newUnnestVariables = ImmutableMap.builder();
    for (Map.Entry<VariableReferenceExpression, List<VariableReferenceExpression>> unnestVariable : node.getUnnestVariables().entrySet()) {
        VariableReferenceExpression input = (VariableReferenceExpression) inlineVariables(context, unnestVariable.getKey());
        ImmutableList.Builder<VariableReferenceExpression> newVariables = ImmutableList.builder();
        for (VariableReferenceExpression variable : unnestVariable.getValue()) {
            VariableReferenceExpression newVariable = variableAllocator.newVariable(variable.getSourceLocation(), "unnest_field", variable.getType());
            context.put(variable, newVariable);
            newVariables.add(newVariable);
        }
        newUnnestVariables.put(input, newVariables.build());
    }
    // Generate canonical ordinality variable
    Optional<VariableReferenceExpression> ordinalityVariable = node.getOrdinalityVariable().map(variable -> {
        VariableReferenceExpression newVariable = variableAllocator.newVariable(variable.getSourceLocation(), "unnest_ordinality", variable.getType());
        context.put(variable, newVariable);
        return newVariable;
    });
    return Optional.of(new UnnestNode(node.getSourceLocation(), planNodeidAllocator.getNextId(), source.get(), node.getReplicateVariables().stream().map(variable -> (VariableReferenceExpression) inlineVariables(context, variable)).collect(toImmutableList()), newUnnestVariables.build(), ordinalityVariable));
}

Example 4 with UnnestNode

use of com.facebook.presto.sql.planner.plan.UnnestNode in project presto by prestodb.

the class RelationPlanner method visitUnnest.

@Override
protected RelationPlan visitUnnest(Unnest node, Void context) {
    Scope scope = analysis.getScope(node);
    ImmutableList.Builder<VariableReferenceExpression> outputVariablesBuilder = ImmutableList.builder();
    for (Field field : scope.getRelationType().getVisibleFields()) {
        VariableReferenceExpression variable = variableAllocator.newVariable(field);
        outputVariablesBuilder.add(variable);
    }
    List<VariableReferenceExpression> unnestedVariables = outputVariablesBuilder.build();
    // If we got here, then we must be unnesting a constant, and not be in a join (where there could be column references)
    ImmutableList.Builder<VariableReferenceExpression> argumentVariables = ImmutableList.builder();
    ImmutableList.Builder<RowExpression> values = ImmutableList.builder();
    ImmutableMap.Builder<VariableReferenceExpression, List<VariableReferenceExpression>> unnestVariables = ImmutableMap.builder();
    Iterator<VariableReferenceExpression> unnestedVariablesIterator = unnestedVariables.iterator();
    for (Expression expression : node.getExpressions()) {
        Type type = analysis.getType(expression);
        Expression rewritten = Coercer.addCoercions(expression, analysis);
        rewritten = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), rewritten);
        values.add(castToRowExpression(rewritten));
        VariableReferenceExpression input = variableAllocator.newVariable(rewritten, type);
        argumentVariables.add(new VariableReferenceExpression(getSourceLocation(rewritten), input.getName(), type));
        if (type instanceof ArrayType) {
            Type elementType = ((ArrayType) type).getElementType();
            if (!SystemSessionProperties.isLegacyUnnest(session) && elementType instanceof RowType) {
                ImmutableList.Builder<VariableReferenceExpression> unnestVariableBuilder = ImmutableList.builder();
                for (int i = 0; i < ((RowType) elementType).getFields().size(); i++) {
                    unnestVariableBuilder.add(unnestedVariablesIterator.next());
                }
                unnestVariables.put(input, unnestVariableBuilder.build());
            } else {
                unnestVariables.put(input, ImmutableList.of(unnestedVariablesIterator.next()));
            }
        } else if (type instanceof MapType) {
            unnestVariables.put(input, ImmutableList.of(unnestedVariablesIterator.next(), unnestedVariablesIterator.next()));
        } else {
            throw new IllegalArgumentException("Unsupported type for UNNEST: " + type);
        }
    }
    Optional<VariableReferenceExpression> ordinalityVariable = node.isWithOrdinality() ? Optional.of(unnestedVariablesIterator.next()) : Optional.empty();
    checkState(!unnestedVariablesIterator.hasNext(), "Not all output variables were matched with input variables");
    ValuesNode valuesNode = new ValuesNode(getSourceLocation(node), idAllocator.getNextId(), argumentVariables.build(), ImmutableList.of(values.build()));
    UnnestNode unnestNode = new UnnestNode(getSourceLocation(node), idAllocator.getNextId(), valuesNode, ImmutableList.of(), unnestVariables.build(), ordinalityVariable);
    return new RelationPlan(unnestNode, scope, unnestedVariables);
}