Search in sources :

Example 26 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class PruneTableScanColumns method pruneColumns.

public static Optional<PlanNode> pruneColumns(Metadata metadata, TypeProvider types, Session session, TableScanNode node, Set<Symbol> referencedOutputs) {
    List<Symbol> newOutputs = filteredCopy(node.getOutputSymbols(), referencedOutputs::contains);
    if (newOutputs.size() == node.getOutputSymbols().size()) {
        return Optional.empty();
    }
    List<ConnectorExpression> projections = newOutputs.stream().map(symbol -> new Variable(symbol.getName(), types.get(symbol))).collect(toImmutableList());
    TableHandle handle = node.getTable();
    Optional<ProjectionApplicationResult<TableHandle>> result = metadata.applyProjection(session, handle, projections, newOutputs.stream().collect(toImmutableMap(Symbol::getName, node.getAssignments()::get)));
    Map<Symbol, ColumnHandle> newAssignments;
    // Bail out if the connector does anything other than limit the list of columns (e.g., if it synthesizes arbitrary expressions)
    if (result.isPresent() && result.get().getProjections().stream().allMatch(Variable.class::isInstance)) {
        handle = result.get().getHandle();
        Map<String, ColumnHandle> assignments = result.get().getAssignments().stream().collect(toImmutableMap(Assignment::getVariable, Assignment::getColumn));
        ImmutableMap.Builder<Symbol, ColumnHandle> builder = ImmutableMap.builder();
        for (int i = 0; i < newOutputs.size(); i++) {
            Variable variable = (Variable) result.get().getProjections().get(i);
            builder.put(newOutputs.get(i), assignments.get(variable.getName()));
        }
        newAssignments = builder.buildOrThrow();
    } else {
        newAssignments = newOutputs.stream().collect(toImmutableMap(Function.identity(), node.getAssignments()::get));
    }
    Set<ColumnHandle> visibleColumns = ImmutableSet.copyOf(newAssignments.values());
    TupleDomain<ColumnHandle> enforcedConstraint = node.getEnforcedConstraint().filter((columnHandle, domain) -> visibleColumns.contains(columnHandle));
    Optional<PlanNodeStatsEstimate> newStatistics = node.getStatistics().map(statistics -> new PlanNodeStatsEstimate(statistics.getOutputRowCount(), statistics.getSymbolStatistics().entrySet().stream().filter(entry -> newAssignments.containsKey(entry.getKey())).collect(toImmutableMap(Entry::getKey, Entry::getValue))));
    return Optional.of(new TableScanNode(node.getId(), handle, newOutputs, newAssignments, enforcedConstraint, newStatistics, node.isUpdateTarget(), node.getUseConnectorNodePartitioning()));
}
Also used : Variable(io.trino.spi.expression.Variable) Function(java.util.function.Function) PlanNode(io.trino.sql.planner.plan.PlanNode) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) TableScanNode(io.trino.sql.planner.plan.TableScanNode) PlanNodeStatsEstimate(io.trino.cost.PlanNodeStatsEstimate) Symbol(io.trino.sql.planner.Symbol) MoreLists.filteredCopy(io.trino.util.MoreLists.filteredCopy) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Set(java.util.Set) TupleDomain(io.trino.spi.predicate.TupleDomain) Patterns.tableScan(io.trino.sql.planner.plan.Patterns.tableScan) List(java.util.List) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) TableHandle(io.trino.metadata.TableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Assignment(io.trino.spi.connector.Assignment) Entry(java.util.Map.Entry) Metadata(io.trino.metadata.Metadata) TypeProvider(io.trino.sql.planner.TypeProvider) Optional(java.util.Optional) Session(io.trino.Session) ColumnHandle(io.trino.spi.connector.ColumnHandle) Variable(io.trino.spi.expression.Variable) PlanNodeStatsEstimate(io.trino.cost.PlanNodeStatsEstimate) Symbol(io.trino.sql.planner.Symbol) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) TableScanNode(io.trino.sql.planner.plan.TableScanNode) TableHandle(io.trino.metadata.TableHandle)

Example 27 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class PushProjectionIntoTableScan method apply.

@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
    TableScanNode tableScan = captures.get(TABLE_SCAN);
    // Extract translatable components from projection expressions. Prepare a mapping from these internal
    // expression nodes to corresponding ConnectorExpression translations.
    Map<NodeRef<Expression>, ConnectorExpression> partialTranslations = project.getAssignments().getMap().entrySet().stream().flatMap(expression -> extractPartialTranslations(expression.getValue(), context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes(), plannerContext).entrySet().stream()).collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (first, ignore) -> first));
    List<NodeRef<Expression>> nodesForPartialProjections = ImmutableList.copyOf(partialTranslations.keySet());
    List<ConnectorExpression> connectorPartialProjections = ImmutableList.copyOf(partialTranslations.values());
    Map<String, Symbol> inputVariableMappings = tableScan.getAssignments().keySet().stream().collect(toImmutableMap(Symbol::getName, identity()));
    Map<String, ColumnHandle> assignments = inputVariableMappings.entrySet().stream().collect(toImmutableMap(Entry::getKey, entry -> tableScan.getAssignments().get(entry.getValue())));
    Optional<ProjectionApplicationResult<TableHandle>> result = plannerContext.getMetadata().applyProjection(context.getSession(), tableScan.getTable(), connectorPartialProjections, assignments);
    if (result.isEmpty()) {
        return Result.empty();
    }
    List<ConnectorExpression> newConnectorPartialProjections = result.get().getProjections();
    checkState(newConnectorPartialProjections.size() == connectorPartialProjections.size(), "Mismatch between input and output projections from the connector: expected %s but got %s", connectorPartialProjections.size(), newConnectorPartialProjections.size());
    List<Symbol> newScanOutputs = new ArrayList<>();
    Map<Symbol, ColumnHandle> newScanAssignments = new HashMap<>();
    Map<String, Symbol> variableMappings = new HashMap<>();
    for (Assignment assignment : result.get().getAssignments()) {
        Symbol symbol = context.getSymbolAllocator().newSymbol(assignment.getVariable(), assignment.getType());
        newScanOutputs.add(symbol);
        newScanAssignments.put(symbol, assignment.getColumn());
        variableMappings.put(assignment.getVariable(), symbol);
    }
    // Translate partial connector projections back to new partial projections
    List<Expression> newPartialProjections = newConnectorPartialProjections.stream().map(expression -> ConnectorExpressionTranslator.translate(context.getSession(), expression, plannerContext, variableMappings, new LiteralEncoder(plannerContext))).collect(toImmutableList());
    // Map internal node references to new partial projections
    ImmutableMap.Builder<NodeRef<Expression>, Expression> nodesToNewPartialProjectionsBuilder = ImmutableMap.builder();
    for (int i = 0; i < nodesForPartialProjections.size(); i++) {
        nodesToNewPartialProjectionsBuilder.put(nodesForPartialProjections.get(i), newPartialProjections.get(i));
    }
    Map<NodeRef<Expression>, Expression> nodesToNewPartialProjections = nodesToNewPartialProjectionsBuilder.buildOrThrow();
    // Stitch partial translations to form new complete projections
    Assignments.Builder newProjectionAssignments = Assignments.builder();
    project.getAssignments().entrySet().forEach(entry -> {
        newProjectionAssignments.put(entry.getKey(), replaceExpression(entry.getValue(), nodesToNewPartialProjections));
    });
    Optional<PlanNodeStatsEstimate> newStatistics = tableScan.getStatistics().map(statistics -> {
        PlanNodeStatsEstimate.Builder builder = PlanNodeStatsEstimate.builder();
        builder.setOutputRowCount(statistics.getOutputRowCount());
        for (int i = 0; i < connectorPartialProjections.size(); i++) {
            ConnectorExpression inputConnectorExpression = connectorPartialProjections.get(i);
            ConnectorExpression resultConnectorExpression = newConnectorPartialProjections.get(i);
            if (!(resultConnectorExpression instanceof Variable)) {
                continue;
            }
            String resultVariableName = ((Variable) resultConnectorExpression).getName();
            Expression inputExpression = ConnectorExpressionTranslator.translate(context.getSession(), inputConnectorExpression, plannerContext, inputVariableMappings, new LiteralEncoder(plannerContext));
            SymbolStatsEstimate symbolStatistics = scalarStatsCalculator.calculate(inputExpression, statistics, context.getSession(), context.getSymbolAllocator().getTypes());
            builder.addSymbolStatistics(variableMappings.get(resultVariableName), symbolStatistics);
        }
        return builder.build();
    });
    verifyTablePartitioning(context, tableScan, result.get().getHandle());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new TableScanNode(tableScan.getId(), result.get().getHandle(), newScanOutputs, newScanAssignments, TupleDomain.all(), newStatistics, tableScan.isUpdateTarget(), tableScan.getUseConnectorNodePartitioning()), newProjectionAssignments.build()));
}
Also used : LiteralEncoder(io.trino.sql.planner.LiteralEncoder) HashMap(java.util.HashMap) Variable(io.trino.spi.expression.Variable) Capture.newCapture(io.trino.matching.Capture.newCapture) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) Verify.verify(com.google.common.base.Verify.verify) NodeRef(io.trino.sql.tree.NodeRef) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) Rule(io.trino.sql.planner.iterative.Rule) ReferenceAwareExpressionNodeInliner.replaceExpression(io.trino.sql.planner.ReferenceAwareExpressionNodeInliner.replaceExpression) ProjectNode(io.trino.sql.planner.plan.ProjectNode) TableScanNode(io.trino.sql.planner.plan.TableScanNode) PlanNodeStatsEstimate(io.trino.cost.PlanNodeStatsEstimate) TablePartitioning(io.trino.metadata.TableProperties.TablePartitioning) Symbol(io.trino.sql.planner.Symbol) SymbolStatsEstimate(io.trino.cost.SymbolStatsEstimate) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) PartialTranslator.extractPartialTranslations(io.trino.sql.planner.PartialTranslator.extractPartialTranslations) ConnectorExpressionTranslator(io.trino.sql.planner.ConnectorExpressionTranslator) Assignments(io.trino.sql.planner.plan.Assignments) TupleDomain(io.trino.spi.predicate.TupleDomain) Patterns.tableScan(io.trino.sql.planner.plan.Patterns.tableScan) ScalarStatsCalculator(io.trino.cost.ScalarStatsCalculator) Capture(io.trino.matching.Capture) Preconditions.checkState(com.google.common.base.Preconditions.checkState) List(java.util.List) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) Pattern(io.trino.matching.Pattern) TableHandle(io.trino.metadata.TableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) TypeAnalyzer(io.trino.sql.planner.TypeAnalyzer) SystemSessionProperties.isAllowPushdownIntoConnectors(io.trino.SystemSessionProperties.isAllowPushdownIntoConnectors) Patterns.source(io.trino.sql.planner.plan.Patterns.source) Captures(io.trino.matching.Captures) Assignment(io.trino.spi.connector.Assignment) Entry(java.util.Map.Entry) Function.identity(java.util.function.Function.identity) Optional(java.util.Optional) Expression(io.trino.sql.tree.Expression) Patterns.project(io.trino.sql.planner.plan.Patterns.project) Session(io.trino.Session) PlannerContext(io.trino.sql.PlannerContext) Variable(io.trino.spi.expression.Variable) HashMap(java.util.HashMap) PlanNodeStatsEstimate(io.trino.cost.PlanNodeStatsEstimate) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Symbol(io.trino.sql.planner.Symbol) ArrayList(java.util.ArrayList) Assignments(io.trino.sql.planner.plan.Assignments) SymbolStatsEstimate(io.trino.cost.SymbolStatsEstimate) Assignment(io.trino.spi.connector.Assignment) NodeRef(io.trino.sql.tree.NodeRef) LiteralEncoder(io.trino.sql.planner.LiteralEncoder) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) TableScanNode(io.trino.sql.planner.plan.TableScanNode) ReferenceAwareExpressionNodeInliner.replaceExpression(io.trino.sql.planner.ReferenceAwareExpressionNodeInliner.replaceExpression) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Expression(io.trino.sql.tree.Expression) ProjectNode(io.trino.sql.planner.plan.ProjectNode) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap)

Example 28 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class DefaultJdbcMetadata method applyAggregation.

@Override
public Optional<AggregationApplicationResult<ConnectorTableHandle>> applyAggregation(ConnectorSession session, ConnectorTableHandle table, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets) {
    if (!isAggregationPushdownEnabled(session)) {
        return Optional.empty();
    }
    JdbcTableHandle handle = (JdbcTableHandle) table;
    // Global aggregation is represented by [[]]
    verify(!groupingSets.isEmpty(), "No grouping sets provided");
    if (!jdbcClient.supportsAggregationPushdown(session, handle, aggregates, assignments, groupingSets)) {
        // JDBC client implementation prevents pushdown for the given table
        return Optional.empty();
    }
    if (handle.getLimit().isPresent()) {
        handle = flushAttributesAsQuery(session, handle);
    }
    int nextSyntheticColumnId = handle.getNextSyntheticColumnId();
    ImmutableList.Builder<JdbcColumnHandle> newColumns = ImmutableList.builder();
    ImmutableList.Builder<ConnectorExpression> projections = ImmutableList.builder();
    ImmutableList.Builder<Assignment> resultAssignments = ImmutableList.builder();
    ImmutableMap.Builder<String, String> expressions = ImmutableMap.builder();
    List<List<JdbcColumnHandle>> groupingSetsAsJdbcColumnHandles = groupingSets.stream().map(groupingSet -> groupingSet.stream().map(JdbcColumnHandle.class::cast).collect(toImmutableList())).collect(toImmutableList());
    Optional<List<JdbcColumnHandle>> tableColumns = handle.getColumns();
    groupingSetsAsJdbcColumnHandles.stream().flatMap(List::stream).distinct().peek(handle.getColumns().<Consumer<JdbcColumnHandle>>map(columns -> groupKey -> verify(columns.contains(groupKey), "applyAggregation called with a grouping column %s which was not included in the table columns: %s", groupKey, tableColumns)).orElse(groupKey -> {
    })).forEach(newColumns::add);
    for (AggregateFunction aggregate : aggregates) {
        Optional<JdbcExpression> expression = jdbcClient.implementAggregation(session, aggregate, assignments);
        if (expression.isEmpty()) {
            return Optional.empty();
        }
        String columnName = SYNTHETIC_COLUMN_NAME_PREFIX + nextSyntheticColumnId;
        nextSyntheticColumnId++;
        JdbcColumnHandle newColumn = JdbcColumnHandle.builder().setColumnName(columnName).setJdbcTypeHandle(expression.get().getJdbcTypeHandle()).setColumnType(aggregate.getOutputType()).setComment(Optional.of("synthetic")).build();
        newColumns.add(newColumn);
        projections.add(new Variable(newColumn.getColumnName(), aggregate.getOutputType()));
        resultAssignments.add(new Assignment(newColumn.getColumnName(), newColumn, aggregate.getOutputType()));
        expressions.put(columnName, expression.get().getExpression());
    }
    List<JdbcColumnHandle> newColumnsList = newColumns.build();
    // We need to have matching column handles in JdbcTableHandle constructed below, as columns read via JDBC must match column handles list.
    // For more context see assertion in JdbcRecordSetProvider.getRecordSet
    PreparedQuery preparedQuery = jdbcClient.prepareQuery(session, handle, Optional.of(groupingSetsAsJdbcColumnHandles), newColumnsList, expressions.buildOrThrow());
    handle = new JdbcTableHandle(new JdbcQueryRelationHandle(preparedQuery), TupleDomain.all(), ImmutableList.of(), Optional.empty(), OptionalLong.empty(), Optional.of(newColumnsList), handle.getAllReferencedTables(), nextSyntheticColumnId);
    return Optional.of(new AggregationApplicationResult<>(handle, projections.build(), resultAssignments.build(), ImmutableMap.of(), false));
}
Also used : SortItem(io.trino.spi.connector.SortItem) AggregateFunction(io.trino.spi.connector.AggregateFunction) TopNApplicationResult(io.trino.spi.connector.TopNApplicationResult) NOT_SUPPORTED(io.trino.spi.StandardErrorCode.NOT_SUPPORTED) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) JdbcMetadataSessionProperties.isAggregationPushdownEnabled(io.trino.plugin.jdbc.JdbcMetadataSessionProperties.isAggregationPushdownEnabled) TableNotFoundException(io.trino.spi.connector.TableNotFoundException) ConnectorOutputTableHandle(io.trino.spi.connector.ConnectorOutputTableHandle) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) Functions.identity(com.google.common.base.Functions.identity) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) Domain(io.trino.spi.predicate.Domain) Collection(java.util.Collection) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) ComputedStatistics(io.trino.spi.statistics.ComputedStatistics) Set(java.util.Set) ConnectorExpressions.extractConjuncts(io.trino.plugin.base.expression.ConnectorExpressions.extractConjuncts) TrinoException(io.trino.spi.TrinoException) LimitApplicationResult(io.trino.spi.connector.LimitApplicationResult) ConnectorOutputMetadata(io.trino.spi.connector.ConnectorOutputMetadata) SchemaTableName(io.trino.spi.connector.SchemaTableName) ConnectorTableSchema(io.trino.spi.connector.ConnectorTableSchema) Preconditions.checkState(com.google.common.base.Preconditions.checkState) List(java.util.List) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) TrinoPrincipal(io.trino.spi.security.TrinoPrincipal) BIGINT(io.trino.spi.type.BigintType.BIGINT) SchemaTablePrefix(io.trino.spi.connector.SchemaTablePrefix) Assignment(io.trino.spi.connector.Assignment) Optional(java.util.Optional) Math.max(java.lang.Math.max) TableScanRedirectApplicationResult(io.trino.spi.connector.TableScanRedirectApplicationResult) SystemTable(io.trino.spi.connector.SystemTable) Types(java.sql.Types) Verify.verifyNotNull(com.google.common.base.Verify.verifyNotNull) AccessDeniedException(io.trino.spi.security.AccessDeniedException) ConnectorTableLayout(io.trino.spi.connector.ConnectorTableLayout) Constraint(io.trino.spi.connector.Constraint) ConnectorInsertTableHandle(io.trino.spi.connector.ConnectorInsertTableHandle) JoinCondition(io.trino.spi.connector.JoinCondition) Slice(io.airlift.slice.Slice) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) ConnectorTableMetadata(io.trino.spi.connector.ConnectorTableMetadata) HashMap(java.util.HashMap) Variable(io.trino.spi.expression.Variable) DomainPushdownResult(io.trino.plugin.jdbc.PredicatePushdownController.DomainPushdownResult) AtomicReference(java.util.concurrent.atomic.AtomicReference) ArrayList(java.util.ArrayList) OptionalLong(java.util.OptionalLong) JoinStatistics(io.trino.spi.connector.JoinStatistics) ImmutableList(com.google.common.collect.ImmutableList) JoinType(io.trino.spi.connector.JoinType) Verify.verify(com.google.common.base.Verify.verify) Objects.requireNonNull(java.util.Objects.requireNonNull) ColumnHandle(io.trino.spi.connector.ColumnHandle) TableStatistics(io.trino.spi.statistics.TableStatistics) JdbcMetadataSessionProperties.isJoinPushdownEnabled(io.trino.plugin.jdbc.JdbcMetadataSessionProperties.isJoinPushdownEnabled) JdbcMetadataSessionProperties.isTopNPushdownEnabled(io.trino.plugin.jdbc.JdbcMetadataSessionProperties.isTopNPushdownEnabled) Constant(io.trino.spi.expression.Constant) ConstraintApplicationResult(io.trino.spi.connector.ConstraintApplicationResult) JdbcMetadataSessionProperties.isComplexExpressionPushdown(io.trino.plugin.jdbc.JdbcMetadataSessionProperties.isComplexExpressionPushdown) ConnectorSession(io.trino.spi.connector.ConnectorSession) TupleDomain(io.trino.spi.predicate.TupleDomain) AggregationApplicationResult(io.trino.spi.connector.AggregationApplicationResult) Consumer(java.util.function.Consumer) JoinApplicationResult(io.trino.spi.connector.JoinApplicationResult) ConnectorExpressions.and(io.trino.plugin.base.expression.ConnectorExpressions.and) ConnectorTableProperties(io.trino.spi.connector.ConnectorTableProperties) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Variable(io.trino.spi.expression.Variable) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) ImmutableList(com.google.common.collect.ImmutableList) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Assignment(io.trino.spi.connector.Assignment) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) List(java.util.List) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) Constraint(io.trino.spi.connector.Constraint) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) AggregateFunction(io.trino.spi.connector.AggregateFunction)

Example 29 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class DefaultJdbcMetadata method getVariableColumnHandle.

private static Optional<JdbcColumnHandle> getVariableColumnHandle(Map<String, ColumnHandle> assignments, ConnectorExpression expression) {
    requireNonNull(assignments, "assignments is null");
    requireNonNull(expression, "expression is null");
    if (!(expression instanceof Variable)) {
        return Optional.empty();
    }
    String name = ((Variable) expression).getName();
    ColumnHandle columnHandle = assignments.get(name);
    verifyNotNull(columnHandle, "No assignment for %s", name);
    return Optional.of(((JdbcColumnHandle) columnHandle));
}
Also used : ColumnHandle(io.trino.spi.connector.ColumnHandle) Variable(io.trino.spi.expression.Variable)

Example 30 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class ImplementAvgFloatingPoint method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(ARGUMENT);
    JdbcColumnHandle columnHandle = (JdbcColumnHandle) context.getAssignment(argument.getName());
    verify(aggregateFunction.getOutputType() == columnHandle.getColumnType());
    return Optional.of(new JdbcExpression(format("avg(%s)", context.rewriteExpression(argument).orElseThrow()), columnHandle.getJdbcTypeHandle()));
}
Also used : Variable(io.trino.spi.expression.Variable) JdbcColumnHandle(io.trino.plugin.jdbc.JdbcColumnHandle) JdbcExpression(io.trino.plugin.jdbc.JdbcExpression)

Aggregations

Variable (io.trino.spi.expression.Variable)51 JdbcExpression (io.trino.plugin.jdbc.JdbcExpression)21 ConnectorExpression (io.trino.spi.expression.ConnectorExpression)21 JdbcColumnHandle (io.trino.plugin.jdbc.JdbcColumnHandle)18 Test (org.testng.annotations.Test)12 ImmutableList (com.google.common.collect.ImmutableList)11 ColumnHandle (io.trino.spi.connector.ColumnHandle)11 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)10 Assignment (io.trino.spi.connector.Assignment)10 Map (java.util.Map)9 Optional (java.util.Optional)9 ImmutableMap (com.google.common.collect.ImmutableMap)8 TupleDomain (io.trino.spi.predicate.TupleDomain)8 List (java.util.List)8 ImmutableMap.toImmutableMap (com.google.common.collect.ImmutableMap.toImmutableMap)7 AggregateFunction (io.trino.spi.connector.AggregateFunction)7 ProjectionApplicationResult (io.trino.spi.connector.ProjectionApplicationResult)7 AggregateExpression (io.trino.plugin.pinot.query.AggregateExpression)6 Domain (io.trino.spi.predicate.Domain)6 Verify.verify (com.google.common.base.Verify.verify)5