Examples with PinotColumnHandle io.trino.plugin.pinot.PinotColumnHandle used on opensource projects

Example 1 with PinotColumnHandle

use of io.trino.plugin.pinot.PinotColumnHandle in project trino by trinodb.

the class DynamicTableBuilder method getPinotColumns.

private static List<PinotColumnHandle> getPinotColumns(SchemaTableName schemaTableName, List<ExpressionContext> expressions, List<String> aliases, Map<String, ColumnHandle> columnHandles, PinotTypeResolver pinotTypeResolver, Map<String, PinotColumnNameAndTrinoType> aggregateTypes) {
    ImmutableList.Builder<PinotColumnHandle> pinotColumnsBuilder = ImmutableList.builder();
    for (int index = 0; index < expressions.size(); index++) {
        ExpressionContext expressionContext = expressions.get(index);
        Optional<String> alias = getAlias(aliases, index);
        // Since Pinot doesn't support subqueries yet, we can only have one occurrence of SELECT *
        if (expressionContext.getType() == ExpressionContext.Type.IDENTIFIER && expressionContext.getIdentifier().equals(WILDCARD)) {
            pinotColumnsBuilder.addAll(columnHandles.values().stream().map(handle -> fromNonAggregateColumnHandle((PinotColumnHandle) handle)).collect(toImmutableList()));
        } else {
            pinotColumnsBuilder.add(getPinotColumnHandle(schemaTableName, expressionContext, alias, columnHandles, pinotTypeResolver, aggregateTypes));
        }
    }
    return pinotColumnsBuilder.build();
}

Example 2 with PinotColumnHandle

use of io.trino.plugin.pinot.PinotColumnHandle in project trino by trinodb.

the class DynamicTableBuilder method buildFromPql.

public static DynamicTable buildFromPql(PinotMetadata pinotMetadata, SchemaTableName schemaTableName, PinotClient pinotClient) {
    requireNonNull(pinotMetadata, "pinotMetadata is null");
    requireNonNull(schemaTableName, "schemaTableName is null");
    String query = schemaTableName.getTableName();
    BrokerRequest request = REQUEST_COMPILER.compileToBrokerRequest(query);
    PinotQuery pinotQuery = request.getPinotQuery();
    QueryContext queryContext = BrokerRequestToQueryContextConverter.convert(request);
    String pinotTableName = stripSuffix(request.getQuerySource().getTableName());
    Optional<String> suffix = getSuffix(request.getQuerySource().getTableName());
    Map<String, ColumnHandle> columnHandles = pinotMetadata.getPinotColumnHandles(pinotTableName);
    List<OrderByExpression> orderBy = ImmutableList.of();
    PinotTypeResolver pinotTypeResolver = new PinotTypeResolver(pinotClient, pinotTableName);
    List<PinotColumnHandle> selectColumns = ImmutableList.of();
    Map<String, PinotColumnNameAndTrinoType> aggregateTypes = ImmutableMap.of();
    if (queryContext.getAggregationFunctions() != null) {
        checkState(queryContext.getAggregationFunctions().length > 0, "Aggregation Functions is empty");
        aggregateTypes = getAggregateTypes(schemaTableName, queryContext, columnHandles);
    }
    if (queryContext.getSelectExpressions() != null) {
        checkState(!queryContext.getSelectExpressions().isEmpty(), "Pinot selections is empty");
        selectColumns = getPinotColumns(schemaTableName, queryContext.getSelectExpressions(), queryContext.getAliasList(), columnHandles, pinotTypeResolver, aggregateTypes);
    }
    if (queryContext.getOrderByExpressions() != null) {
        ImmutableList.Builder<OrderByExpression> orderByBuilder = ImmutableList.builder();
        for (OrderByExpressionContext orderByExpressionContext : queryContext.getOrderByExpressions()) {
            ExpressionContext expressionContext = orderByExpressionContext.getExpression();
            PinotColumnHandle pinotColumnHandle = getPinotColumnHandle(schemaTableName, expressionContext, Optional.empty(), columnHandles, pinotTypeResolver, aggregateTypes);
            orderByBuilder.add(new OrderByExpression(pinotColumnHandle.getExpression(), orderByExpressionContext.isAsc()));
        }
        orderBy = orderByBuilder.build();
    }
    List<PinotColumnHandle> groupByColumns = ImmutableList.of();
    if (queryContext.getGroupByExpressions() != null) {
        groupByColumns = getPinotColumns(schemaTableName, queryContext.getGroupByExpressions(), ImmutableList.of(), columnHandles, pinotTypeResolver, aggregateTypes);
    }
    Optional<String> filter = Optional.empty();
    if (pinotQuery.getFilterExpression() != null) {
        String formatted = formatFilter(schemaTableName, queryContext.getFilter(), columnHandles);
        filter = Optional.of(formatted);
    }
    return new DynamicTable(pinotTableName, suffix, selectColumns, filter, groupByColumns, ImmutableList.of(), orderBy, OptionalLong.of(queryContext.getLimit()), getOffset(queryContext), query);
}

Example 3 with PinotColumnHandle

use of io.trino.plugin.pinot.PinotColumnHandle in project trino by trinodb.

the class DynamicTableBuilder method getPinotColumnHandle.

private static PinotColumnHandle getPinotColumnHandle(SchemaTableName schemaTableName, ExpressionContext expressionContext, Optional<String> alias, Map<String, ColumnHandle> columnHandles, PinotTypeResolver pinotTypeResolver, Map<String, PinotColumnNameAndTrinoType> aggregateTypes) {
    ExpressionContext rewritten = rewriteExpression(schemaTableName, expressionContext, columnHandles);
    // If there is no alias, pinot autogenerates the column name:
    String columnName = rewritten.toString();
    String pinotExpression = formatExpression(schemaTableName, rewritten);
    Type trinoType;
    boolean isAggregate = hasAggregate(rewritten);
    if (isAggregate) {
        trinoType = requireNonNull(aggregateTypes.get(columnName).getTrinoType(), format("Unexpected aggregate expression: '%s'", rewritten));
        // For aggregation queries, the column name is set by the schema returned from PostAggregationHandler, see getAggregateTypes
        columnName = aggregateTypes.get(columnName).getPinotColumnName();
    } else {
        trinoType = getTrinoTypeFromPinotType(pinotTypeResolver.resolveExpressionType(rewritten, schemaTableName, columnHandles));
        if (!aggregateTypes.isEmpty() && trinoType instanceof ArrayType) {
            trinoType = ((ArrayType) trinoType).getElementType();
        }
    }
    return new PinotColumnHandle(alias.orElse(columnName), trinoType, pinotExpression, alias.isPresent(), isAggregate, isReturnNullOnEmptyGroup(expressionContext), Optional.empty(), Optional.empty());
}

Example 4 with PinotColumnHandle

use of io.trino.plugin.pinot.PinotColumnHandle in project trino by trinodb.

the class ImplementAvg method rewrite.

@Override
public Optional<AggregateExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<Void> context) {
    Variable argument = captures.get(ARGUMENT);
    PinotColumnHandle columnHandle = (PinotColumnHandle) context.getAssignment(argument.getName());
    return Optional.of(new AggregateExpression(aggregateFunction.getFunctionName(), identifierQuote.apply(columnHandle.getColumnName()), true));
}

Example 5 with PinotColumnHandle

use of io.trino.plugin.pinot.PinotColumnHandle in project trino by trinodb.

the class ImplementMinMax method rewrite.

@Override
public Optional<AggregateExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<Void> context) {
    Variable argument = captures.get(ARGUMENT);
    PinotColumnHandle columnHandle = (PinotColumnHandle) context.getAssignment(argument.getName());
    verify(columnHandle.getDataType().equals(aggregateFunction.getOutputType()));
    return Optional.of(new AggregateExpression(aggregateFunction.getFunctionName(), identifierQuote.apply(columnHandle.getColumnName()), true));
}