Examples with DynamicTable io.trino.plugin.pinot.query.DynamicTable used on opensource projects

Example 16 with DynamicTable

use of io.trino.plugin.pinot.query.DynamicTable in project trino by trinodb.

the class TestDynamicTable method testRealtimeOnlyDynamicTable.

@Test
public void testRealtimeOnlyDynamicTable() {
    String tableName = hybridTable.getTableName();
    String tableNameWithSuffix = tableName + REALTIME_SUFFIX;
    String query = format("select * from %s limit 70", tableNameWithSuffix);
    DynamicTable dynamicTable = buildFromPql(pinotMetadata, new SchemaTableName("default", query), mockClusterInfoFetcher);
    String expectedPql = format("select %s from %s limit 70", getColumnNames(tableName).stream().map(TestDynamicTable::quoteIdentifier).collect(joining(", ")), tableNameWithSuffix);
    assertEquals(extractPql(dynamicTable, TupleDomain.all(), ImmutableList.of()), expectedPql);
    assertEquals(dynamicTable.getTableName(), tableName);
}

Example 17 with DynamicTable

use of io.trino.plugin.pinot.query.DynamicTable in project trino by trinodb.

the class PinotMetadata method getTableMetadata.

@Override
public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle table) {
    PinotTableHandle pinotTableHandle = (PinotTableHandle) table;
    if (pinotTableHandle.getQuery().isPresent()) {
        DynamicTable dynamicTable = pinotTableHandle.getQuery().get();
        ImmutableList.Builder<ColumnMetadata> columnMetadataBuilder = ImmutableList.builder();
        for (PinotColumnHandle pinotColumnHandle : dynamicTable.getProjections()) {
            columnMetadataBuilder.add(pinotColumnHandle.getColumnMetadata());
        }
        dynamicTable.getAggregateColumns().forEach(columnHandle -> columnMetadataBuilder.add(columnHandle.getColumnMetadata()));
        SchemaTableName schemaTableName = new SchemaTableName(pinotTableHandle.getSchemaName(), dynamicTable.getTableName());
        return new ConnectorTableMetadata(schemaTableName, columnMetadataBuilder.build());
    }
    SchemaTableName tableName = new SchemaTableName(pinotTableHandle.getSchemaName(), pinotTableHandle.getTableName());
    return getTableMetadata(tableName);
}

Example 18 with DynamicTable

use of io.trino.plugin.pinot.query.DynamicTable in project trino by trinodb.

the class PinotMetadata method applyAggregation.

@Override
public Optional<AggregationApplicationResult<ConnectorTableHandle>> applyAggregation(ConnectorSession session, ConnectorTableHandle handle, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets) {
    if (!isAggregationPushdownEnabled(session)) {
        return Optional.empty();
    }
    // Global aggregation is represented by [[]]
    verify(!groupingSets.isEmpty(), "No grouping sets provided");
    // Pinot currently only supports simple GROUP BY clauses with a single grouping set
    if (groupingSets.size() != 1) {
        return Optional.empty();
    }
    // See https://github.com/apache/pinot/issues/8353 for more details.
    if (getOnlyElement(groupingSets).stream().filter(columnHandle -> ((PinotColumnHandle) columnHandle).getDataType() instanceof ArrayType).findFirst().isPresent()) {
        return Optional.empty();
    }
    PinotTableHandle tableHandle = (PinotTableHandle) handle;
    // If there is an offset then do not push the aggregation down as the results will not be correct
    if (tableHandle.getQuery().isPresent() && (!tableHandle.getQuery().get().getAggregateColumns().isEmpty() || tableHandle.getQuery().get().isAggregateInProjections() || tableHandle.getQuery().get().getOffset().isPresent())) {
        return Optional.empty();
    }
    ImmutableList.Builder<ConnectorExpression> projections = ImmutableList.builder();
    ImmutableList.Builder<Assignment> resultAssignments = ImmutableList.builder();
    ImmutableList.Builder<PinotColumnHandle> aggregateColumnsBuilder = ImmutableList.builder();
    for (AggregateFunction aggregate : aggregates) {
        Optional<AggregateExpression> rewriteResult = aggregateFunctionRewriter.rewrite(session, aggregate, assignments);
        rewriteResult = applyCountDistinct(session, aggregate, assignments, tableHandle, rewriteResult);
        if (rewriteResult.isEmpty()) {
            return Optional.empty();
        }
        AggregateExpression aggregateExpression = rewriteResult.get();
        PinotColumnHandle pinotColumnHandle = new PinotColumnHandle(aggregateExpression.toFieldName(), aggregate.getOutputType(), aggregateExpression.toExpression(), false, true, aggregateExpression.isReturnNullOnEmptyGroup(), Optional.of(aggregateExpression.getFunction()), Optional.of(aggregateExpression.getArgument()));
        aggregateColumnsBuilder.add(pinotColumnHandle);
        projections.add(new Variable(pinotColumnHandle.getColumnName(), pinotColumnHandle.getDataType()));
        resultAssignments.add(new Assignment(pinotColumnHandle.getColumnName(), pinotColumnHandle, pinotColumnHandle.getDataType()));
    }
    List<PinotColumnHandle> groupingColumns = getOnlyElement(groupingSets).stream().map(PinotColumnHandle.class::cast).map(PinotColumnHandle::fromNonAggregateColumnHandle).collect(toImmutableList());
    OptionalLong limitForDynamicTable = OptionalLong.empty();
    // know when the limit was exceeded and throw an error
    if (tableHandle.getLimit().isEmpty() && !groupingColumns.isEmpty()) {
        limitForDynamicTable = OptionalLong.of(maxRowsPerBrokerQuery + 1);
    }
    List<PinotColumnHandle> aggregationColumns = aggregateColumnsBuilder.build();
    String newQuery = "";
    List<PinotColumnHandle> newSelections = groupingColumns;
    if (tableHandle.getQuery().isPresent()) {
        newQuery = tableHandle.getQuery().get().getQuery();
        Map<String, PinotColumnHandle> projectionsMap = tableHandle.getQuery().get().getProjections().stream().collect(toImmutableMap(PinotColumnHandle::getColumnName, identity()));
        groupingColumns = groupingColumns.stream().map(groupIngColumn -> projectionsMap.getOrDefault(groupIngColumn.getColumnName(), groupIngColumn)).collect(toImmutableList());
        ImmutableList.Builder<PinotColumnHandle> newSelectionsBuilder = ImmutableList.<PinotColumnHandle>builder().addAll(groupingColumns);
        aggregationColumns = aggregationColumns.stream().map(aggregateExpression -> resolveAggregateExpressionWithAlias(aggregateExpression, projectionsMap)).collect(toImmutableList());
        newSelections = newSelectionsBuilder.build();
    }
    DynamicTable dynamicTable = new DynamicTable(tableHandle.getTableName(), Optional.empty(), newSelections, tableHandle.getQuery().flatMap(DynamicTable::getFilter), groupingColumns, aggregationColumns, ImmutableList.of(), limitForDynamicTable, OptionalLong.empty(), newQuery);
    tableHandle = new PinotTableHandle(tableHandle.getSchemaName(), tableHandle.getTableName(), tableHandle.getConstraint(), tableHandle.getLimit(), Optional.of(dynamicTable));
    return Optional.of(new AggregationApplicationResult<>(tableHandle, projections.build(), resultAssignments.build(), ImmutableMap.of(), false));
}

Example 19 with DynamicTable

use of io.trino.plugin.pinot.query.DynamicTable in project trino by trinodb.

the class TestDynamicTable method testOrderBy.

@Test
public void testOrderBy() {
    String tableName = hybridTable.getTableName();
    String tableNameWithSuffix = tableName + REALTIME_SUFFIX;
    String query = format("select ArrDelay + 34 - DaysSinceEpoch, FlightNum from %s order by ArrDelay asc, DaysSinceEpoch desc", tableNameWithSuffix);
    DynamicTable dynamicTable = buildFromPql(pinotMetadata, new SchemaTableName("default", query), mockClusterInfoFetcher);
    String expectedPql = format("select plus(\"ArrDelay\", '34') - \"DaysSinceEpoch\", \"FlightNum\" from %s order by \"ArrDelay\", \"DaysSinceEpoch\" desc limit 10", tableNameWithSuffix);
    assertEquals(extractPql(dynamicTable, TupleDomain.all(), ImmutableList.of()), expectedPql);
    assertEquals(dynamicTable.getTableName(), tableName);
}

Example 20 with DynamicTable

use of io.trino.plugin.pinot.query.DynamicTable in project trino by trinodb.

the class PinotPageSourceProvider method createPageSource.

@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle tableHandle, List<ColumnHandle> columns, DynamicFilter dynamicFilter) {
    requireNonNull(split, "split is null");
    PinotSplit pinotSplit = (PinotSplit) split;
    List<PinotColumnHandle> handles = new ArrayList<>();
    for (ColumnHandle handle : columns) {
        handles.add((PinotColumnHandle) handle);
    }
    PinotTableHandle pinotTableHandle = (PinotTableHandle) tableHandle;
    String query = generatePql(pinotTableHandle, handles, pinotSplit.getSuffix(), pinotSplit.getTimePredicate(), limitForSegmentQueries);
    switch(pinotSplit.getSplitType()) {
        case SEGMENT:
            return new PinotSegmentPageSource(session, estimatedNonNumericColumnSize, limitForSegmentQueries, this.pinotQueryClient, pinotSplit, handles, query);
        case BROKER:
            PinotQueryInfo pinotQueryInfo;
            if (pinotTableHandle.getQuery().isPresent()) {
                DynamicTable dynamicTable = pinotTableHandle.getQuery().get();
                pinotQueryInfo = new PinotQueryInfo(dynamicTable.getTableName(), extractPql(dynamicTable, pinotTableHandle.getConstraint(), handles), dynamicTable.getGroupingColumns().size());
            } else {
                pinotQueryInfo = new PinotQueryInfo(pinotTableHandle.getTableName(), query, 0);
            }
            return new PinotBrokerPageSource(session, pinotQueryInfo, handles, clusterInfoFetcher, limitForBrokerQueries);
    }
    throw new UnsupportedOperationException("Unknown Pinot split type: " + pinotSplit.getSplitType());
}