private static Optional<Symbol> newRadiusSymbol(Context context, Expression expression) { if (expression instanceof SymbolReference) { return Optional.empty(); } return Optional.of(context.getSymbolAllocator().newSymbol(expression, DOUBLE)); }
private static TableWriterNode rewriteSource( TableWriterNode writerNode, UnionNode unionNode, int source, List<Map<Symbol, Symbol>> sourceMappings, Context context) { Map<Symbol, Symbol> inputMappings = getInputSymbolMapping(unionNode, source); ImmutableMap.Builder<Symbol, Symbol> mappings = ImmutableMap.builder(); mappings.putAll(inputMappings); ImmutableMap.Builder<Symbol, Symbol> outputMappings = ImmutableMap.builder(); for (Symbol outputSymbol : writerNode.getOutputSymbols()) { if (inputMappings.containsKey(outputSymbol)) { outputMappings.put(outputSymbol, inputMappings.get(outputSymbol)); } else { Symbol newSymbol = context.getSymbolAllocator().newSymbol(outputSymbol); outputMappings.put(outputSymbol, newSymbol); mappings.put(outputSymbol, newSymbol); } } sourceMappings.add(outputMappings.build()); SymbolMapper symbolMapper = new SymbolMapper(mappings.build()); return symbolMapper.map(writerNode, unionNode.getSources().get(source), context.getIdAllocator().getNextId()); }
private static Optional<Symbol> newGeometrySymbol(Context context, Expression expression, Metadata metadata) { if (expression instanceof SymbolReference) { return Optional.empty(); } return Optional.of(context.getSymbolAllocator().newSymbol(expression, metadata.getType(GEOMETRY_TYPE_SIGNATURE))); }
private PlanNodeWithCost getSemiJoinNodeWithCost(SemiJoinNode possibleJoinNode, Context context) TypeProvider types = context.getSymbolAllocator().getTypes(); StatsProvider stats = context.getStatsProvider(); boolean replicated = possibleJoinNode.getDistributionType().get().equals(REPLICATED);
Type type = context.getSymbolAllocator().getTypes().get(entry.getKey()); Symbol symbol = context.getSymbolAllocator().newSymbol(translatedExpression, type); assignments.put(symbol, translatedExpression); projectSymbolMapping.put(entry.getKey(), symbol);
private PlanNodeWithCost getJoinNodeWithCost(Context context, JoinNode possibleJoinNode) TypeProvider types = context.getSymbolAllocator().getTypes(); StatsProvider stats = context.getStatsProvider(); boolean replicated = possibleJoinNode.getDistributionType().get().equals(REPLICATED);
private Optional<PlanNode> rewriteToNonDefaultAggregation(ApplyNode applyNode, Context context) { checkState(applyNode.getSubquery().getOutputSymbols().isEmpty(), "Expected subquery output symbols to be pruned"); Symbol exists = getOnlyElement(applyNode.getSubqueryAssignments().getSymbols()); Symbol subqueryTrue = context.getSymbolAllocator().newSymbol("subqueryTrue", BOOLEAN); Assignments.Builder assignments = Assignments.builder(); assignments.putIdentities(applyNode.getInput().getOutputSymbols()); assignments.put(exists, new CoalesceExpression(ImmutableList.of(subqueryTrue.toSymbolReference(), BooleanLiteral.FALSE_LITERAL))); PlanNode subquery = new ProjectNode( context.getIdAllocator().getNextId(), new LimitNode( context.getIdAllocator().getNextId(), applyNode.getSubquery(), 1L, false), Assignments.of(subqueryTrue, TRUE_LITERAL)); PlanNodeDecorrelator decorrelator = new PlanNodeDecorrelator(context.getIdAllocator(), context.getLookup()); if (!decorrelator.decorrelateFilters(subquery, applyNode.getCorrelation()).isPresent()) { return Optional.empty(); } return Optional.of(new ProjectNode(context.getIdAllocator().getNextId(), new LateralJoinNode( applyNode.getId(), applyNode.getInput(), subquery, applyNode.getCorrelation(), LEFT, applyNode.getOriginSubquery()), assignments.build())); }
private PlanNode rewriteToDefaultAggregation(ApplyNode parent, Context context) { Symbol count = context.getSymbolAllocator().newSymbol(COUNT.toString(), BIGINT); Symbol exists = getOnlyElement(parent.getSubqueryAssignments().getSymbols()); return new LateralJoinNode( parent.getId(), parent.getInput(), new ProjectNode( context.getIdAllocator().getNextId(), new AggregationNode( context.getIdAllocator().getNextId(), parent.getSubquery(), ImmutableMap.of(count, new Aggregation(COUNT_CALL, countSignature, Optional.empty())), globalAggregation(), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), Optional.empty()), Assignments.of(exists, new ComparisonExpression(GREATER_THAN, count.toSymbolReference(), new Cast(new LongLiteral("0"), BIGINT.toString())))), parent.getCorrelation(), INNER, parent.getOriginSubquery()); } }
@Override public Result apply(LateralJoinNode lateralJoinNode, Captures captures, Context context) { PlanNode subquery = lateralJoinNode.getSubquery(); if (!isScalar(subquery, context.getLookup())) { return Result.empty(); } Optional<AggregationNode> aggregation = findAggregation(subquery, context.getLookup()); if (!(aggregation.isPresent() && aggregation.get().getGroupingKeys().isEmpty())) { return Result.empty(); } ScalarAggregationToJoinRewriter rewriter = new ScalarAggregationToJoinRewriter(functionRegistry, context.getSymbolAllocator(), context.getIdAllocator(), context.getLookup()); PlanNode rewrittenNode = rewriter.rewriteScalarAggregation(lateralJoinNode, aggregation.get()); if (rewrittenNode instanceof LateralJoinNode) { return Result.empty(); } return Result.ofPlanNode(rewrittenNode); }
private static PlanNode addPartitioningNodes(Context context, PlanNode node, Symbol partitionSymbol, KdbTree kdbTree, Expression geometry, Optional<Expression> radius) { Assignments.Builder projections = Assignments.builder(); for (Symbol outputSymbol : node.getOutputSymbols()) { projections.putIdentity(outputSymbol); } ImmutableList.Builder<Expression> partitioningArguments = ImmutableList.<Expression>builder() .add(new Cast(new StringLiteral(KdbTreeUtils.toJson(kdbTree)), KDB_TREE_TYPENAME)) .add(geometry); radius.map(partitioningArguments::add); FunctionCall partitioningFunction = new FunctionCall(QualifiedName.of("spatial_partitions"), partitioningArguments.build()); Symbol partitionsSymbol = context.getSymbolAllocator().newSymbol(partitioningFunction, new ArrayType(INTEGER)); projections.put(partitionsSymbol, partitioningFunction); return new UnnestNode( context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), node, projections.build()), node.getOutputSymbols(), ImmutableMap.of(partitionsSymbol, ImmutableList.of(partitionSymbol)), Optional.empty()); }
public static boolean canReplicate(JoinNode joinNode, Context context) { JoinDistributionType joinDistributionType = getJoinDistributionType(context.getSession()); if (!joinDistributionType.canReplicate()) { return false; } Optional<DataSize> joinMaxBroadcastTableSize = getJoinMaxBroadcastTableSize(context.getSession()); if (!joinMaxBroadcastTableSize.isPresent()) { return true; } PlanNode buildSide = joinNode.getRight(); PlanNodeStatsEstimate buildSideStatsEstimate = context.getStatsProvider().getStats(buildSide); double buildSideSizeInBytes = buildSideStatsEstimate.getOutputSizeInBytes(buildSide.getOutputSymbols(), context.getSymbolAllocator().getTypes()); return buildSideSizeInBytes <= joinMaxBroadcastTableSize.get().toBytes(); }
@Override public Result apply(ApplyNode apply, Captures captures, Context context) { Assignments subqueryAssignments = apply.getSubqueryAssignments(); if (subqueryAssignments.size() != 1) { return Result.empty(); } Expression assignmentExpression = getOnlyElement(subqueryAssignments.getExpressions()); if (!(assignmentExpression instanceof InPredicate)) { return Result.empty(); } InPredicate inPredicate = (InPredicate) assignmentExpression; Symbol inPredicateOutputSymbol = getOnlyElement(subqueryAssignments.getSymbols()); return apply(apply, inPredicate, inPredicateOutputSymbol, context.getLookup(), context.getIdAllocator(), context.getSymbolAllocator()); }
private static <T> RuleApplication applyRule(Rule<T> rule, PlanNode planNode, Rule.Context context) { PlanNodeMatcher matcher = new PlanNodeMatcher(context.getLookup()); Match<T> match = matcher.match(rule.getPattern(), planNode); Rule.Result result; if (!rule.isEnabled(context.getSession()) || match.isEmpty()) { result = Rule.Result.empty(); } else { result = rule.apply(match.value(), match.captures(), context); } return new RuleApplication(context.getLookup(), context.getStatsProvider(), context.getSymbolAllocator().getTypes(), result); }
@Override public Result apply(FilterNode filterNode, Captures captures, Context context) { TableScanNode tableScan = captures.get(TABLE_SCAN); PlanNode rewritten = planTableScan(tableScan, filterNode.getPredicate(), context.getSession(), context.getSymbolAllocator().getTypes(), context.getIdAllocator(), metadata, parser, domainTranslator); if (arePlansSame(filterNode, tableScan, rewritten)) { return Result.empty(); } return Result.ofPlanNode(rewritten); }
private static ExpressionRewriter createRewrite(Metadata metadata, SqlParser sqlParser) { requireNonNull(metadata, "metadata is null"); requireNonNull(sqlParser, "sqlParser is null"); return (expression, context) -> DesugarAtTimeZoneRewriter.rewrite(expression, context.getSession(), metadata, sqlParser, context.getSymbolAllocator()); } }
private boolean canReplicate(SemiJoinNode node, Context context) { Optional<DataSize> joinMaxBroadcastTableSize = getJoinMaxBroadcastTableSize(context.getSession()); if (!joinMaxBroadcastTableSize.isPresent()) { return true; } PlanNode buildSide = node.getFilteringSource(); PlanNodeStatsEstimate buildSideStatsEstimate = context.getStatsProvider().getStats(buildSide); double buildSideSizeInBytes = buildSideStatsEstimate.getOutputSizeInBytes(buildSide.getOutputSymbols(), context.getSymbolAllocator().getTypes()); return buildSideSizeInBytes <= joinMaxBroadcastTableSize.get().toBytes(); }
@Override public Result apply(TableScanNode tableScanNode, Captures captures, Context context) { if (tableScanNode.getLayout().isPresent()) { return Result.empty(); } return Result.ofPlanNode(planTableScan(tableScanNode, TRUE_LITERAL, context.getSession(), context.getSymbolAllocator().getTypes(), context.getIdAllocator(), metadata, parser, domainTranslator)); } }
private static Expression rewrite(Expression expression, Rule.Context context) { return LambdaCaptureDesugaringRewriter.rewrite(expression, context.getSymbolAllocator().getTypes(), context.getSymbolAllocator()); } }
private static Type getExpressionType(Expression expression, Context context, Metadata metadata, SqlParser sqlParser) { Type type = getExpressionTypes(context.getSession(), metadata, sqlParser, context.getSymbolAllocator().getTypes(), expression, emptyList(), WarningCollector.NOOP) .get(NodeRef.of(expression)); verify(type != null); return type; }
private static ExpressionRewriter createRewrite(Metadata metadata, SqlParser sqlParser) { requireNonNull(metadata, "metadata is null"); requireNonNull(sqlParser, "sqlParser is null"); LiteralEncoder literalEncoder = new LiteralEncoder(metadata.getBlockEncodingSerde()); return (expression, context) -> rewrite(expression, context.getSession(), context.getSymbolAllocator(), metadata, literalEncoder, sqlParser); } }