Examples with Node io.trino.sql.tree.Node used on opensource projects

Example 1 with Node

use of io.trino.sql.tree.Node in project trino by trinodb.

the class QueryPlanner method planGroupingSets.

private GroupingSetsPlan planGroupingSets(PlanBuilder subPlan, QuerySpecification node, GroupingSetAnalysis groupingSetAnalysis) {
    Map<Symbol, Symbol> groupingSetMappings = new LinkedHashMap<>();
    // Compute a set of artificial columns that will contain the values of the original columns
    // filtered by whether the column is included in the grouping set
    // This will become the basis for the scope for any column references
    Symbol[] fields = new Symbol[subPlan.getTranslations().getFieldSymbols().size()];
    for (FieldId field : groupingSetAnalysis.getAllFields()) {
        Symbol input = subPlan.getTranslations().getFieldSymbols().get(field.getFieldIndex());
        Symbol output = symbolAllocator.newSymbol(input, "gid");
        fields[field.getFieldIndex()] = output;
        groupingSetMappings.put(output, input);
    }
    Map<ScopeAware<Expression>, Symbol> complexExpressions = new HashMap<>();
    for (Expression expression : groupingSetAnalysis.getComplexExpressions()) {
        if (!complexExpressions.containsKey(scopeAwareKey(expression, analysis, subPlan.getScope()))) {
            Symbol input = subPlan.translate(expression);
            Symbol output = symbolAllocator.newSymbol(expression, analysis.getType(expression), "gid");
            complexExpressions.put(scopeAwareKey(expression, analysis, subPlan.getScope()), output);
            groupingSetMappings.put(output, input);
        }
    }
    // For the purpose of "distinct", we need to canonicalize column references that may have varying
    // syntactic forms (e.g., "t.a" vs "a"). Thus we need to enumerate grouping sets based on the underlying
    // fieldId associated with each column reference expression.
    // The catch is that simple group-by expressions can be arbitrary expressions (this is a departure from the SQL specification).
    // But, they don't affect the number of grouping sets or the behavior of "distinct" . We can compute all the candidate
    // grouping sets in terms of fieldId, dedup as appropriate and then cross-join them with the complex expressions.
    // This tracks the grouping sets before complex expressions are considered.
    // It's also used to compute the descriptors needed to implement grouping()
    List<Set<FieldId>> columnOnlyGroupingSets = enumerateGroupingSets(groupingSetAnalysis);
    if (node.getGroupBy().isPresent() && node.getGroupBy().get().isDistinct()) {
        columnOnlyGroupingSets = columnOnlyGroupingSets.stream().distinct().collect(toImmutableList());
    }
    // translate from FieldIds to Symbols
    List<List<Symbol>> sets = columnOnlyGroupingSets.stream().map(set -> set.stream().map(FieldId::getFieldIndex).map(index -> fields[index]).collect(toImmutableList())).collect(toImmutableList());
    // combine (cartesian product) with complex expressions
    List<List<Symbol>> groupingSets = sets.stream().map(set -> ImmutableList.<Symbol>builder().addAll(set).addAll(complexExpressions.values()).build()).collect(toImmutableList());
    // Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
    PlanNode groupId;
    Optional<Symbol> groupIdSymbol = Optional.empty();
    if (groupingSets.size() > 1) {
        groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
        groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSets, groupingSetMappings, subPlan.getRoot().getOutputSymbols(), groupIdSymbol.get());
    } else {
        Assignments.Builder assignments = Assignments.builder();
        assignments.putIdentities(subPlan.getRoot().getOutputSymbols());
        groupingSetMappings.forEach((key, value) -> assignments.put(key, value.toSymbolReference()));
        groupId = new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), assignments.build());
    }
    subPlan = new PlanBuilder(subPlan.getTranslations().withNewMappings(complexExpressions, Arrays.asList(fields)), groupId);
    return new GroupingSetsPlan(subPlan, columnOnlyGroupingSets, groupingSets, groupIdSymbol);
}

Example 2 with Node

use of io.trino.sql.tree.Node in project trino by trinodb.

the class QueryPlanner method plan.

public UpdateNode plan(Update node) {
    Table table = node.getTable();
    TableHandle handle = analysis.getTableHandle(table);
    TableSchema tableSchema = plannerContext.getMetadata().getTableSchema(session, handle);
    Map<String, ColumnHandle> columnMap = plannerContext.getMetadata().getColumnHandles(session, handle);
    List<ColumnSchema> columnsSchemas = tableSchema.getColumns();
    List<String> targetColumnNames = node.getAssignments().stream().map(assignment -> assignment.getName().getValue()).collect(toImmutableList());
    // Create lists of columnnames and SET expressions, in table column order
    ImmutableList.Builder<String> updatedColumnNamesBuilder = ImmutableList.builder();
    ImmutableList.Builder<ColumnHandle> updatedColumnHandlesBuilder = ImmutableList.builder();
    ImmutableList.Builder<Expression> orderedColumnValuesBuilder = ImmutableList.builder();
    for (ColumnSchema columnSchema : columnsSchemas) {
        String name = columnSchema.getName();
        int index = targetColumnNames.indexOf(name);
        if (index >= 0) {
            updatedColumnNamesBuilder.add(name);
            updatedColumnHandlesBuilder.add(requireNonNull(columnMap.get(name), "columnMap didn't contain name"));
            orderedColumnValuesBuilder.add(node.getAssignments().get(index).getValue());
        }
    }
    List<String> updatedColumnNames = updatedColumnNamesBuilder.build();
    List<ColumnHandle> updatedColumnHandles = updatedColumnHandlesBuilder.build();
    List<Expression> orderedColumnValues = orderedColumnValuesBuilder.build();
    // create table scan
    RelationPlan relationPlan = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, recursiveSubqueries).process(table, null);
    PlanBuilder builder = newPlanBuilder(relationPlan, analysis, lambdaDeclarationToSymbolMap);
    if (node.getWhere().isPresent()) {
        builder = filter(builder, node.getWhere().get(), node);
    }
    builder = subqueryPlanner.handleSubqueries(builder, orderedColumnValues, analysis.getSubqueries(node));
    builder = builder.appendProjections(orderedColumnValues, symbolAllocator, idAllocator);
    PlanAndMappings planAndMappings = coerce(builder, orderedColumnValues, analysis, idAllocator, symbolAllocator, typeCoercion);
    builder = planAndMappings.getSubPlan();
    ImmutableList.Builder<Symbol> updatedColumnValuesBuilder = ImmutableList.builder();
    orderedColumnValues.forEach(columnValue -> updatedColumnValuesBuilder.add(planAndMappings.get(columnValue)));
    Symbol rowId = builder.translate(analysis.getRowIdField(table));
    updatedColumnValuesBuilder.add(rowId);
    List<Symbol> outputs = ImmutableList.of(symbolAllocator.newSymbol("partialrows", BIGINT), symbolAllocator.newSymbol("fragment", VARBINARY));
    Optional<PlanNodeId> tableScanId = getIdForLeftTableScan(relationPlan.getRoot());
    checkArgument(tableScanId.isPresent(), "tableScanId not present");
    // create update node
    return new UpdateNode(idAllocator.getNextId(), builder.getRoot(), new UpdateTarget(Optional.empty(), plannerContext.getMetadata().getTableMetadata(session, handle).getTable(), updatedColumnNames, updatedColumnHandles), rowId, updatedColumnValuesBuilder.build(), outputs);
}

Example 3 with Node

use of io.trino.sql.tree.Node in project trino by trinodb.

the class RelationPlanner method visitValues.

@Override
protected RelationPlan visitValues(Values node, Void context) {
    Scope scope = analysis.getScope(node);
    ImmutableList.Builder<Symbol> outputSymbolsBuilder = ImmutableList.builder();
    for (Field field : scope.getRelationType().getVisibleFields()) {
        Symbol symbol = symbolAllocator.newSymbol(field);
        outputSymbolsBuilder.add(symbol);
    }
    List<Symbol> outputSymbols = outputSymbolsBuilder.build();
    TranslationMap translationMap = new TranslationMap(outerContext, analysis.getScope(node), analysis, lambdaDeclarationToSymbolMap, outputSymbols);
    ImmutableList.Builder<Expression> rows = ImmutableList.builder();
    for (Expression row : node.getRows()) {
        if (row instanceof Row) {
            rows.add(new Row(((Row) row).getItems().stream().map(item -> coerceIfNecessary(analysis, item, translationMap.rewrite(item))).collect(toImmutableList())));
        } else if (analysis.getType(row) instanceof RowType) {
            rows.add(coerceIfNecessary(analysis, row, translationMap.rewrite(row)));
        } else {
            rows.add(new Row(ImmutableList.of(coerceIfNecessary(analysis, row, translationMap.rewrite(row)))));
        }
    }
    ValuesNode valuesNode = new ValuesNode(idAllocator.getNextId(), outputSymbols, rows.build());
    return new RelationPlan(valuesNode, scope, outputSymbols, outerContext);
}

Example 4 with Node

use of io.trino.sql.tree.Node in project trino by trinodb.

the class AstBuilder method visitQueryNoWith.

@Override
public Node visitQueryNoWith(SqlBaseParser.QueryNoWithContext context) {
    QueryBody term = (QueryBody) visit(context.queryTerm());
    Optional<OrderBy> orderBy = Optional.empty();
    if (context.ORDER() != null) {
        orderBy = Optional.of(new OrderBy(getLocation(context.ORDER()), visit(context.sortItem(), SortItem.class)));
    }
    Optional<Offset> offset = Optional.empty();
    if (context.OFFSET() != null) {
        Expression rowCount;
        if (context.offset.INTEGER_VALUE() != null) {
            rowCount = new LongLiteral(getLocation(context.offset.INTEGER_VALUE()), context.offset.getText());
        } else {
            rowCount = new Parameter(getLocation(context.offset.QUESTION_MARK()), parameterPosition);
            parameterPosition++;
        }
        offset = Optional.of(new Offset(Optional.of(getLocation(context.OFFSET())), rowCount));
    }
    Optional<Node> limit = Optional.empty();
    if (context.FETCH() != null) {
        Optional<Expression> rowCount = Optional.empty();
        if (context.fetchFirst != null) {
            if (context.fetchFirst.INTEGER_VALUE() != null) {
                rowCount = Optional.of(new LongLiteral(getLocation(context.fetchFirst.INTEGER_VALUE()), context.fetchFirst.getText()));
            } else {
                rowCount = Optional.of(new Parameter(getLocation(context.fetchFirst.QUESTION_MARK()), parameterPosition));
                parameterPosition++;
            }
        }
        limit = Optional.of(new FetchFirst(Optional.of(getLocation(context.FETCH())), rowCount, context.TIES() != null));
    } else if (context.LIMIT() != null) {
        if (context.limit == null) {
            throw new IllegalStateException("Missing LIMIT value");
        }
        Expression rowCount;
        if (context.limit.ALL() != null) {
            rowCount = new AllRows(getLocation(context.limit.ALL()));
        } else if (context.limit.rowCount().INTEGER_VALUE() != null) {
            rowCount = new LongLiteral(getLocation(context.limit.rowCount().INTEGER_VALUE()), context.limit.getText());
        } else {
            rowCount = new Parameter(getLocation(context.limit.rowCount().QUESTION_MARK()), parameterPosition);
            parameterPosition++;
        }
        limit = Optional.of(new Limit(Optional.of(getLocation(context.LIMIT())), rowCount));
    }
    if (term instanceof QuerySpecification) {
        // When we have a simple query specification
        // followed by order by, offset, limit or fetch,
        // fold the order by, limit, offset or fetch clauses
        // into the query specification (analyzer/planner
        // expects this structure to resolve references with respect
        // to columns defined in the query specification)
        QuerySpecification query = (QuerySpecification) term;
        return new Query(getLocation(context), Optional.empty(), new QuerySpecification(getLocation(context), query.getSelect(), query.getFrom(), query.getWhere(), query.getGroupBy(), query.getHaving(), query.getWindows(), orderBy, offset, limit), Optional.empty(), Optional.empty(), Optional.empty());
    }
    return new Query(getLocation(context), Optional.empty(), term, orderBy, offset, limit);
}

Example 5 with Node

use of io.trino.sql.tree.Node in project trino by trinodb.

the class AstUtils method treeHash.

/**
 * <p>Computes a hash of the given AST by applying the provided subtree hasher at each level.</p>
 *
 * <p>If the hasher returns a non-empty {@link OptionalInt}, the value is treated as the hash for
 * the subtree at that node. Otherwise, the hashes of its children are computed and combined.</p>
 */
public static int treeHash(Node node, Function<Node, OptionalInt> subtreeHasher) {
    OptionalInt hash = subtreeHasher.apply(node);
    if (hash.isPresent()) {
        return hash.getAsInt();
    }
    List<? extends Node> children = node.getChildren();
    int result = node.getClass().hashCode();
    for (Node element : children) {
        result = 31 * result + treeHash(element, subtreeHasher);
    }
    return result;
}