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

Example 1 with QuerySpecification

use of io.trino.sql.tree.QuerySpecification 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 QuerySpecification

use of io.trino.sql.tree.QuerySpecification 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 3 with QuerySpecification

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

the class TestSqlParser method testShowStatsForQuery.

@Test
public void testShowStatsForQuery() {
    String[] tableNames = { "t", "s.t", "c.s.t" };
    for (String fullName : tableNames) {
        QualifiedName qualifiedName = makeQualifiedName(fullName);
        // Simple SELECT
        assertStatement(format("SHOW STATS FOR (SELECT * FROM %s)", qualifiedName), createShowStats(qualifiedName, ImmutableList.of(new AllColumns()), Optional.empty()));
        // SELECT with predicate
        assertStatement(format("SHOW STATS FOR (SELECT * FROM %s WHERE field > 0)", qualifiedName), createShowStats(qualifiedName, ImmutableList.of(new AllColumns()), Optional.of(new ComparisonExpression(ComparisonExpression.Operator.GREATER_THAN, new Identifier("field"), new LongLiteral("0")))));
        // SELECT with more complex predicate
        assertStatement(format("SHOW STATS FOR (SELECT * FROM %s WHERE field > 0 or field < 0)", qualifiedName), createShowStats(qualifiedName, ImmutableList.of(new AllColumns()), Optional.of(LogicalExpression.or(new ComparisonExpression(ComparisonExpression.Operator.GREATER_THAN, new Identifier("field"), new LongLiteral("0")), new ComparisonExpression(ComparisonExpression.Operator.LESS_THAN, new Identifier("field"), new LongLiteral("0"))))));
    }
    // SELECT with LIMIT
    assertThat(statement("SHOW STATS FOR (SELECT * FROM t LIMIT 10)")).isEqualTo(new ShowStats(Optional.of(location(1, 1)), new TableSubquery(new Query(location(1, 17), Optional.empty(), new QuerySpecification(location(1, 17), new Select(location(1, 17), false, ImmutableList.of(new AllColumns(location(1, 24), Optional.empty(), ImmutableList.of()))), Optional.of(new Table(location(1, 31), QualifiedName.of(ImmutableList.of(new Identifier(location(1, 31), "t", false))))), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.of(new Limit(location(1, 33), new LongLiteral(location(1, 39), "10")))), Optional.empty(), Optional.empty(), Optional.empty()))));
    // SELECT with ORDER BY ... LIMIT
    assertThat(statement("SHOW STATS FOR (SELECT * FROM t ORDER BY field LIMIT 10)")).isEqualTo(new ShowStats(Optional.of(location(1, 1)), new TableSubquery(new Query(location(1, 17), Optional.empty(), new QuerySpecification(location(1, 17), new Select(location(1, 17), false, ImmutableList.of(new AllColumns(location(1, 24), Optional.empty(), ImmutableList.of()))), Optional.of(new Table(location(1, 31), QualifiedName.of(ImmutableList.of(new Identifier(location(1, 31), "t", false))))), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.of(new OrderBy(location(1, 33), ImmutableList.of(new SortItem(location(1, 42), new Identifier(location(1, 42), "field", false), ASCENDING, UNDEFINED)))), Optional.empty(), Optional.of(new Limit(location(1, 48), new LongLiteral(location(1, 54), "10")))), Optional.empty(), Optional.empty(), Optional.empty()))));
    // SELECT with WITH
    assertThat(statement("SHOW STATS FOR (\n" + "   WITH t AS (SELECT 1 )\n" + "   SELECT * FROM t)")).isEqualTo(new ShowStats(Optional.of(location(1, 1)), new TableSubquery(new Query(location(2, 4), Optional.of(new With(location(2, 4), false, ImmutableList.of(new WithQuery(location(2, 9), new Identifier(location(2, 9), "t", false), new Query(location(2, 15), Optional.empty(), new QuerySpecification(location(2, 15), new Select(location(2, 15), false, ImmutableList.of(new SingleColumn(location(2, 22), new LongLiteral(location(2, 22), "1"), Optional.empty()))), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty())))), new QuerySpecification(location(3, 4), new Select(location(3, 4), false, ImmutableList.of(new AllColumns(location(3, 11), Optional.empty(), ImmutableList.of()))), Optional.of(new Table(location(3, 18), QualifiedName.of(ImmutableList.of(new Identifier(location(3, 18), "t", false))))), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty(), Optional.empty(), Optional.empty()))));
}

Example 4 with QuerySpecification

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

the class QueryUtil method emptyQuery.

// TODO pass column types
public static Query emptyQuery(List<String> columns) {
    Select select = selectList(columns.stream().map(column -> new SingleColumn(new NullLiteral(), QueryUtil.identifier(column))).toArray(SelectItem[]::new));
    Optional<Expression> where = Optional.of(FALSE_LITERAL);
    return query(new QuerySpecification(select, Optional.empty(), where, Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty()));
}

Example 5 with QuerySpecification

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

the class TestSqlParser method testQueryPeriod.

@Test
public void testQueryPeriod() {
    Expression rangeValue = new TimestampLiteral(location(1, 37), "2021-03-01 00:00:01");
    QueryPeriod queryPeriod = new QueryPeriod(location(1, 17), QueryPeriod.RangeType.TIMESTAMP, rangeValue);
    Table table = new Table(location(1, 15), qualifiedName(location(1, 15), "t"), queryPeriod);
    assertThat(statement("SELECT * FROM t FOR TIMESTAMP AS OF TIMESTAMP '2021-03-01 00:00:01'")).isEqualTo(new Query(location(1, 1), Optional.empty(), new QuerySpecification(location(1, 1), new Select(location(1, 1), false, ImmutableList.of(new AllColumns(location(1, 8), Optional.empty(), ImmutableList.of()))), Optional.of(table), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty(), Optional.empty(), Optional.empty()));
    rangeValue = new StringLiteral(location(1, 35), "version1");
    queryPeriod = new QueryPeriod(new NodeLocation(1, 17), QueryPeriod.RangeType.VERSION, rangeValue);
    table = new Table(location(1, 15), qualifiedName(location(1, 15), "t"), queryPeriod);
    assertThat(statement("SELECT * FROM t FOR VERSION AS OF 'version1'")).isEqualTo(new Query(location(1, 1), Optional.empty(), new QuerySpecification(location(1, 1), new Select(location(1, 1), false, ImmutableList.of(new AllColumns(location(1, 8), Optional.empty(), ImmutableList.of()))), Optional.of(table), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty(), Optional.empty(), Optional.empty()));
}