Examples with Context com.facebook.presto.sql.planner.iterative.Rule.Context used on opensource projects

Example 1 with Context

use of com.facebook.presto.sql.planner.iterative.Rule.Context in project presto by prestodb.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<VariableReferenceExpression> outputVariables, CallExpression spatialComparison, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
    FunctionMetadata spatialComparisonMetadata = metadata.getFunctionAndTypeManager().getFunctionMetadata(spatialComparison.getFunctionHandle());
    checkArgument(spatialComparison.getArguments().size() == 2 && spatialComparisonMetadata.getOperatorType().isPresent());
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    List<VariableReferenceExpression> leftVariables = leftNode.getOutputVariables();
    List<VariableReferenceExpression> rightVariables = rightNode.getOutputVariables();
    RowExpression radius;
    Optional<VariableReferenceExpression> newRadiusVariable;
    CallExpression newComparison;
    if (spatialComparisonMetadata.getOperatorType().get() == OperatorType.LESS_THAN || spatialComparisonMetadata.getOperatorType().get() == OperatorType.LESS_THAN_OR_EQUAL) {
        // ST_Distance(a, b) <= r
        radius = spatialComparison.getArguments().get(1);
        Set<VariableReferenceExpression> radiusVariables = extractUnique(radius);
        if (radiusVariables.isEmpty() || (rightVariables.containsAll(radiusVariables) && containsNone(leftVariables, radiusVariables))) {
            newRadiusVariable = newRadiusVariable(context, radius);
            newComparison = new CallExpression(spatialComparison.getSourceLocation(), spatialComparison.getDisplayName(), spatialComparison.getFunctionHandle(), spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(0), mapToExpression(newRadiusVariable, radius)));
        } else {
            return Result.empty();
        }
    } else {
        // r >= ST_Distance(a, b)
        radius = spatialComparison.getArguments().get(0);
        Set<VariableReferenceExpression> radiusVariables = extractUnique(radius);
        if (radiusVariables.isEmpty() || (rightVariables.containsAll(radiusVariables) && containsNone(leftVariables, radiusVariables))) {
            newRadiusVariable = newRadiusVariable(context, radius);
            OperatorType flippedOperatorType = flip(spatialComparisonMetadata.getOperatorType().get());
            FunctionHandle flippedHandle = getFlippedFunctionHandle(spatialComparison, metadata.getFunctionAndTypeManager());
            newComparison = new CallExpression(spatialComparison.getSourceLocation(), // TODO verify if this is the correct function displayName
            flippedOperatorType.getOperator(), flippedHandle, spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(1), mapToExpression(newRadiusVariable, radius)));
        } else {
            return Result.empty();
        }
    }
    RowExpression newFilter = replaceExpression(filter, ImmutableMap.of(spatialComparison, newComparison));
    PlanNode newRightNode = newRadiusVariable.map(variable -> addProjection(context, rightNode, variable, radius)).orElse(rightNode);
    JoinNode newJoinNode = new JoinNode(joinNode.getSourceLocation(), joinNode.getId(), joinNode.getType(), leftNode, newRightNode, joinNode.getCriteria(), joinNode.getOutputVariables(), Optional.of(newFilter), joinNode.getLeftHashVariable(), joinNode.getRightHashVariable(), joinNode.getDistributionType(), joinNode.getDynamicFilters());
    return tryCreateSpatialJoin(context, newJoinNode, newFilter, nodeId, outputVariables, (CallExpression) newComparison.getArguments().get(0), Optional.of(newComparison.getArguments().get(1)), metadata, splitManager, pageSourceManager);
}

Example 2 with Context

use of com.facebook.presto.sql.planner.iterative.Rule.Context in project presto by prestodb.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<VariableReferenceExpression> outputVariables, CallExpression spatialFunction, Optional<RowExpression> radius, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
    FunctionAndTypeManager functionAndTypeManager = metadata.getFunctionAndTypeManager();
    List<RowExpression> arguments = spatialFunction.getArguments();
    verify(arguments.size() == 2);
    RowExpression firstArgument = arguments.get(0);
    RowExpression secondArgument = arguments.get(1);
    // Currently, only inner joins are supported for spherical geometries.
    if (joinNode.getType() != INNER && isSphericalJoin(metadata, firstArgument, secondArgument)) {
        return Result.empty();
    }
    Set<VariableReferenceExpression> firstVariables = extractUnique(firstArgument);
    Set<VariableReferenceExpression> secondVariables = extractUnique(secondArgument);
    if (firstVariables.isEmpty() || secondVariables.isEmpty()) {
        return Result.empty();
    }
    // If either firstArgument or secondArgument is not a
    // VariableReferenceExpression, will replace the left/right join node
    // with a projection that adds the argument as a variable.
    Optional<VariableReferenceExpression> newFirstVariable = newGeometryVariable(context, firstArgument);
    Optional<VariableReferenceExpression> newSecondVariable = newGeometryVariable(context, secondArgument);
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    PlanNode newLeftNode;
    PlanNode newRightNode;
    // Check if the order of arguments of the spatial function matches the order of join sides
    int alignment = checkAlignment(joinNode, firstVariables, secondVariables);
    if (alignment > 0) {
        newLeftNode = newFirstVariable.map(variable -> addProjection(context, leftNode, variable, firstArgument)).orElse(leftNode);
        newRightNode = newSecondVariable.map(variable -> addProjection(context, rightNode, variable, secondArgument)).orElse(rightNode);
    } else if (alignment < 0) {
        newLeftNode = newSecondVariable.map(variable -> addProjection(context, leftNode, variable, secondArgument)).orElse(leftNode);
        newRightNode = newFirstVariable.map(variable -> addProjection(context, rightNode, variable, firstArgument)).orElse(rightNode);
    } else {
        return Result.empty();
    }
    RowExpression newFirstArgument = mapToExpression(newFirstVariable, firstArgument);
    RowExpression newSecondArgument = mapToExpression(newSecondVariable, secondArgument);
    // Implement partitioned spatial joins:
    // If the session parameter points to a valid spatial partitioning, use
    // that to assign to each probe and build rows the partitions that the
    // geometry intersects.  This is a projection that adds an array of ints
    // which is subsequently unnested.
    Optional<String> spatialPartitioningTableName = canPartitionSpatialJoin(joinNode) ? getSpatialPartitioningTableName(context.getSession()) : Optional.empty();
    Optional<KdbTree> kdbTree = spatialPartitioningTableName.map(tableName -> loadKdbTree(tableName, context.getSession(), metadata, splitManager, pageSourceManager));
    Optional<VariableReferenceExpression> leftPartitionVariable = Optional.empty();
    Optional<VariableReferenceExpression> rightPartitionVariable = Optional.empty();
    if (kdbTree.isPresent()) {
        leftPartitionVariable = Optional.of(context.getVariableAllocator().newVariable(newFirstArgument.getSourceLocation(), "pid", INTEGER));
        rightPartitionVariable = Optional.of(context.getVariableAllocator().newVariable(newSecondArgument.getSourceLocation(), "pid", INTEGER));
        if (alignment > 0) {
            newLeftNode = addPartitioningNodes(context, functionAndTypeManager, newLeftNode, leftPartitionVariable.get(), kdbTree.get(), newFirstArgument, Optional.empty());
            newRightNode = addPartitioningNodes(context, functionAndTypeManager, newRightNode, rightPartitionVariable.get(), kdbTree.get(), newSecondArgument, radius);
        } else {
            newLeftNode = addPartitioningNodes(context, functionAndTypeManager, newLeftNode, leftPartitionVariable.get(), kdbTree.get(), newSecondArgument, Optional.empty());
            newRightNode = addPartitioningNodes(context, functionAndTypeManager, newRightNode, rightPartitionVariable.get(), kdbTree.get(), newFirstArgument, radius);
        }
    }
    CallExpression newSpatialFunction = new CallExpression(spatialFunction.getSourceLocation(), spatialFunction.getDisplayName(), spatialFunction.getFunctionHandle(), spatialFunction.getType(), ImmutableList.of(newFirstArgument, newSecondArgument));
    RowExpression newFilter = RowExpressionNodeInliner.replaceExpression(filter, ImmutableMap.of(spatialFunction, newSpatialFunction));
    return Result.ofPlanNode(new SpatialJoinNode(joinNode.getSourceLocation(), nodeId, SpatialJoinNode.Type.fromJoinNodeType(joinNode.getType()), newLeftNode, newRightNode, outputVariables, newFilter, leftPartitionVariable, rightPartitionVariable, kdbTree.map(KdbTreeUtils::toJson)));
}