Search in sources :

Example 6 with ExpressionContext

use of org.apache.pinot.common.request.context.ExpressionContext in project trino by trinodb.

the class DynamicTableBuilder method getAggregateTypes.

private static Map<String, PinotColumnNameAndTrinoType> getAggregateTypes(SchemaTableName schemaTableName, QueryContext queryContext, Map<String, ColumnHandle> columnHandles) {
    // A mapping from pinot expression to the returned pinot column name and trino type
    // Note: the column name is set by the PostAggregationHandler
    List<ExpressionContext> aggregateColumnExpressions = queryContext.getSelectExpressions().stream().filter(DynamicTableBuilder::hasAggregate).collect(toImmutableList());
    queryContext = new QueryContext.Builder().setAliasList(queryContext.getAliasList()).setSelectExpressions(aggregateColumnExpressions).build();
    DataSchema preAggregationSchema = getPreAggregationDataSchema(queryContext);
    PostAggregationHandler postAggregationHandler = new PostAggregationHandler(queryContext, preAggregationSchema);
    DataSchema postAggregtionSchema = postAggregationHandler.getResultDataSchema();
    ImmutableMap.Builder<String, PinotColumnNameAndTrinoType> aggregationTypesBuilder = ImmutableMap.builder();
    for (int index = 0; index < postAggregtionSchema.size(); index++) {
        aggregationTypesBuilder.put(// Quoting of identifiers is not done to match the corresponding column name in the ResultTable returned from Pinot. Quoting will be done by `DynamicTablePqlExtractor`.
        rewriteExpression(schemaTableName, aggregateColumnExpressions.get(index), columnHandles).toString(), new PinotColumnNameAndTrinoType(postAggregtionSchema.getColumnName(index), toTrinoType(postAggregtionSchema.getColumnDataType(index))));
    }
    return aggregationTypesBuilder.buildOrThrow();
}
Also used : DataSchema(org.apache.pinot.common.utils.DataSchema) ExpressionContext(org.apache.pinot.common.request.context.ExpressionContext) OrderByExpressionContext(org.apache.pinot.common.request.context.OrderByExpressionContext) PostAggregationHandler(org.apache.pinot.core.query.reduce.PostAggregationHandler) ImmutableMap(com.google.common.collect.ImmutableMap)

Aggregations

ExpressionContext (org.apache.pinot.common.request.context.ExpressionContext)6 OrderByExpressionContext (org.apache.pinot.common.request.context.OrderByExpressionContext)4 PinotColumnHandle (io.trino.plugin.pinot.PinotColumnHandle)3 ImmutableList (com.google.common.collect.ImmutableList)2 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)2 ColumnHandle (io.trino.spi.connector.ColumnHandle)2 FunctionContext (org.apache.pinot.common.request.context.FunctionContext)2 ImmutableMap (com.google.common.collect.ImmutableMap)1 Maps.immutableEnumMap (com.google.common.collect.Maps.immutableEnumMap)1 PinotColumnHandle.fromNonAggregateColumnHandle (io.trino.plugin.pinot.PinotColumnHandle.fromNonAggregateColumnHandle)1 PinotColumnHandle.getTrinoTypeFromPinotType (io.trino.plugin.pinot.PinotColumnHandle.getTrinoTypeFromPinotType)1 PinotPatterns.functionContext (io.trino.plugin.pinot.query.PinotPatterns.functionContext)1 PinotSqlFormatter.getColumnHandle (io.trino.plugin.pinot.query.PinotSqlFormatter.getColumnHandle)1 ArrayType (io.trino.spi.type.ArrayType)1 Type (io.trino.spi.type.Type)1 HashMap (java.util.HashMap)1 Map (java.util.Map)1 BrokerRequest (org.apache.pinot.common.request.BrokerRequest)1 PinotQuery (org.apache.pinot.common.request.PinotQuery)1 FilterContext (org.apache.pinot.common.request.context.FilterContext)1