Examples with Result io.prestosql.sql.planner.iterative.Rule.Result used on opensource projects

Example 1 with Result

use of io.prestosql.sql.planner.iterative.Rule.Result in project hetu-core by openlookeng.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<Symbol> outputSymbols, CallExpression spatialComparison, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager, TypeAnalyzer typeAnalyzer) {
    FunctionMetadata spatialComparisonMetadata = metadata.getFunctionAndTypeManager().getFunctionMetadata(spatialComparison.getFunctionHandle());
    checkArgument(spatialComparison.getArguments().size() == 2 && spatialComparisonMetadata.getOperatorType().isPresent());
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    List<Symbol> leftSymbols = leftNode.getOutputSymbols();
    List<Symbol> rightSymbols = rightNode.getOutputSymbols();
    RowExpression radius;
    Optional<Symbol> newRadiusSymbol;
    CallExpression newComparison;
    OperatorType operatorType = spatialComparisonMetadata.getOperatorType().get();
    if (operatorType.equals(OperatorType.LESS_THAN) || operatorType.equals(OperatorType.LESS_THAN_OR_EQUAL)) {
        // ST_Distance(a, b) <= r
        radius = spatialComparison.getArguments().get(1);
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new CallExpression(spatialComparison.getDisplayName(), spatialComparison.getFunctionHandle(), spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(0), mapToExpression(newRadiusSymbol, radius, context)), Optional.empty());
        } else {
            return Result.empty();
        }
    } else {
        // r >= ST_Distance(a, b)
        radius = spatialComparison.getArguments().get(0);
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            OperatorType newOperatorType = SpatialJoinUtils.flip(operatorType);
            FunctionHandle flippedHandle = getFlippedFunctionHandle(spatialComparison, metadata.getFunctionAndTypeManager());
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new CallExpression(newOperatorType.name(), flippedHandle, spatialComparison.getType(), ImmutableList.of(spatialComparison.getArguments().get(1), mapToExpression(newRadiusSymbol, radius, context)), Optional.empty());
        } else {
            return Result.empty();
        }
    }
    RowExpression newFilter = replaceExpression(filter, ImmutableMap.of(spatialComparison, newComparison));
    PlanNode newRightNode = newRadiusSymbol.map(symbol -> addProjection(context, rightNode, symbol, radius)).orElse(rightNode);
    JoinNode newJoinNode = new JoinNode(joinNode.getId(), joinNode.getType(), leftNode, newRightNode, joinNode.getCriteria(), joinNode.getOutputSymbols(), Optional.of(newFilter), joinNode.getLeftHashSymbol(), joinNode.getRightHashSymbol(), joinNode.getDistributionType(), joinNode.isSpillable(), joinNode.getDynamicFilters());
    return tryCreateSpatialJoin(context, newJoinNode, newFilter, nodeId, outputSymbols, (CallExpression) newComparison.getArguments().get(0), Optional.of(newComparison.getArguments().get(1)), metadata, splitManager, pageSourceManager, typeAnalyzer);
}

Example 2 with Result

use of io.prestosql.sql.planner.iterative.Rule.Result in project hetu-core by openlookeng.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, RowExpression filter, PlanNodeId nodeId, List<Symbol> outputSymbols, CallExpression spatialFunction, Optional<RowExpression> radius, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager, TypeAnalyzer typeAnalyzer) {
    // TODO Add support for distributed left spatial joins
    Optional<String> spatialPartitioningTableName = joinNode.getType() == INNER ? getSpatialPartitioningTableName(context.getSession()) : Optional.empty();
    Optional<KdbTree> kdbTree = spatialPartitioningTableName.map(tableName -> loadKdbTree(tableName, context.getSession(), metadata, splitManager, pageSourceManager, nodeId));
    List<RowExpression> arguments = spatialFunction.getArguments();
    verify(arguments.size() == 2);
    RowExpression firstArgument = arguments.get(0);
    RowExpression secondArgument = arguments.get(1);
    Type sphericalGeographyType = metadata.getType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE);
    if (firstArgument.getType().equals(sphericalGeographyType) || secondArgument.getType().equals(sphericalGeographyType)) {
        if (joinNode.getType() != INNER) {
            return Result.empty();
        }
    }
    Set<Symbol> firstSymbols = extractUnique(firstArgument);
    Set<Symbol> secondSymbols = extractUnique(secondArgument);
    if (firstSymbols.isEmpty() || secondSymbols.isEmpty()) {
        return Result.empty();
    }
    Optional<Symbol> newFirstSymbol = newGeometrySymbol(context, firstArgument);
    Optional<Symbol> newSecondSymbol = newGeometrySymbol(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, firstSymbols, secondSymbols);
    if (alignment > 0) {
        newLeftNode = newFirstSymbol.map(symbol -> addProjection(context, leftNode, symbol, firstArgument)).orElse(leftNode);
        newRightNode = newSecondSymbol.map(symbol -> addProjection(context, rightNode, symbol, secondArgument)).orElse(rightNode);
    } else if (alignment < 0) {
        newLeftNode = newSecondSymbol.map(symbol -> addProjection(context, leftNode, symbol, secondArgument)).orElse(leftNode);
        newRightNode = newFirstSymbol.map(symbol -> addProjection(context, rightNode, symbol, firstArgument)).orElse(rightNode);
    } else {
        return Result.empty();
    }
    RowExpression newFirstArgument = mapToExpression(newFirstSymbol, firstArgument, context);
    RowExpression newSecondArgument = mapToExpression(newSecondSymbol, secondArgument, context);
    Optional<Symbol> leftPartitionSymbol = Optional.empty();
    Optional<Symbol> rightPartitionSymbol = Optional.empty();
    if (kdbTree.isPresent()) {
        leftPartitionSymbol = Optional.of(context.getSymbolAllocator().newSymbol("pid", INTEGER));
        rightPartitionSymbol = Optional.of(context.getSymbolAllocator().newSymbol("pid", INTEGER));
        if (alignment > 0) {
            newLeftNode = addPartitioningNodes(metadata, context, newLeftNode, leftPartitionSymbol.get(), kdbTree.get(), newFirstArgument, Optional.empty());
            newRightNode = addPartitioningNodes(metadata, context, newRightNode, rightPartitionSymbol.get(), kdbTree.get(), newSecondArgument, radius);
        } else {
            newLeftNode = addPartitioningNodes(metadata, context, newLeftNode, leftPartitionSymbol.get(), kdbTree.get(), newSecondArgument, Optional.empty());
            newRightNode = addPartitioningNodes(metadata, context, newRightNode, rightPartitionSymbol.get(), kdbTree.get(), newFirstArgument, radius);
        }
    }
    CallExpression newSpatialFunction = new CallExpression(spatialFunction.getDisplayName(), spatialFunction.getFunctionHandle(), spatialFunction.getType(), ImmutableList.of(newFirstArgument, newSecondArgument), Optional.empty());
    RowExpression newFilter = replaceExpression(filter, ImmutableMap.of(spatialFunction, newSpatialFunction));
    return Result.ofPlanNode(new SpatialJoinNode(nodeId, SpatialJoinNode.Type.fromJoinNodeType(joinNode.getType()), newLeftNode, newRightNode, outputSymbols, newFilter, leftPartitionSymbol, rightPartitionSymbol, kdbTree.map(KdbTreeUtils::toJson)));
}