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);
}
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);
}
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);
}
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);
}
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;
}
Aggregations