Examples with ColumnSchema io.trino.spi.connector.ColumnSchema used on opensource projects

Example 1 with ColumnSchema

use of io.trino.spi.connector.ColumnSchema in project trino by trinodb.

the class QueryPlanner method plan.

public UpdateNode plan(Update node) {
    Table table = node.getTable();
    TableHandle handle = analysis.getTableHandle(table);
    TableSchema tableSchema = plannerContext.getMetadata().getTableSchema(session, handle);
    Map<String, ColumnHandle> columnMap = plannerContext.getMetadata().getColumnHandles(session, handle);
    List<ColumnSchema> columnsSchemas = tableSchema.getColumns();
    List<String> targetColumnNames = node.getAssignments().stream().map(assignment -> assignment.getName().getValue()).collect(toImmutableList());
    // Create lists of columnnames and SET expressions, in table column order
    ImmutableList.Builder<String> updatedColumnNamesBuilder = ImmutableList.builder();
    ImmutableList.Builder<ColumnHandle> updatedColumnHandlesBuilder = ImmutableList.builder();
    ImmutableList.Builder<Expression> orderedColumnValuesBuilder = ImmutableList.builder();
    for (ColumnSchema columnSchema : columnsSchemas) {
        String name = columnSchema.getName();
        int index = targetColumnNames.indexOf(name);
        if (index >= 0) {
            updatedColumnNamesBuilder.add(name);
            updatedColumnHandlesBuilder.add(requireNonNull(columnMap.get(name), "columnMap didn't contain name"));
            orderedColumnValuesBuilder.add(node.getAssignments().get(index).getValue());
        }
    }
    List<String> updatedColumnNames = updatedColumnNamesBuilder.build();
    List<ColumnHandle> updatedColumnHandles = updatedColumnHandlesBuilder.build();
    List<Expression> orderedColumnValues = orderedColumnValuesBuilder.build();
    // create table scan
    RelationPlan relationPlan = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, recursiveSubqueries).process(table, null);
    PlanBuilder builder = newPlanBuilder(relationPlan, analysis, lambdaDeclarationToSymbolMap);
    if (node.getWhere().isPresent()) {
        builder = filter(builder, node.getWhere().get(), node);
    }
    builder = subqueryPlanner.handleSubqueries(builder, orderedColumnValues, analysis.getSubqueries(node));
    builder = builder.appendProjections(orderedColumnValues, symbolAllocator, idAllocator);
    PlanAndMappings planAndMappings = coerce(builder, orderedColumnValues, analysis, idAllocator, symbolAllocator, typeCoercion);
    builder = planAndMappings.getSubPlan();
    ImmutableList.Builder<Symbol> updatedColumnValuesBuilder = ImmutableList.builder();
    orderedColumnValues.forEach(columnValue -> updatedColumnValuesBuilder.add(planAndMappings.get(columnValue)));
    Symbol rowId = builder.translate(analysis.getRowIdField(table));
    updatedColumnValuesBuilder.add(rowId);
    List<Symbol> outputs = ImmutableList.of(symbolAllocator.newSymbol("partialrows", BIGINT), symbolAllocator.newSymbol("fragment", VARBINARY));
    Optional<PlanNodeId> tableScanId = getIdForLeftTableScan(relationPlan.getRoot());
    checkArgument(tableScanId.isPresent(), "tableScanId not present");
    // create update node
    return new UpdateNode(idAllocator.getNextId(), builder.getRoot(), new UpdateTarget(Optional.empty(), plannerContext.getMetadata().getTableMetadata(session, handle).getTable(), updatedColumnNames, updatedColumnHandles), rowId, updatedColumnValuesBuilder.build(), outputs);
}