Examples with UpdateTarget io.trino.sql.planner.plan.TableWriterNode.UpdateTarget used on opensource projects

Example 1 with UpdateTarget

use of io.trino.sql.planner.plan.TableWriterNode.UpdateTarget in project trino by trinodb.

the class LocalExecutionPlanner method createTableFinisher.

private static TableFinisher createTableFinisher(Session session, TableFinishNode node, Metadata metadata) {
    WriterTarget target = node.getTarget();
    return (fragments, statistics, tableExecuteContext) -> {
        if (target instanceof CreateTarget) {
            return metadata.finishCreateTable(session, ((CreateTarget) target).getHandle(), fragments, statistics);
        } else if (target instanceof InsertTarget) {
            return metadata.finishInsert(session, ((InsertTarget) target).getHandle(), fragments, statistics);
        } else if (target instanceof TableWriterNode.RefreshMaterializedViewTarget) {
            TableWriterNode.RefreshMaterializedViewTarget refreshTarget = (TableWriterNode.RefreshMaterializedViewTarget) target;
            return metadata.finishRefreshMaterializedView(session, refreshTarget.getTableHandle(), refreshTarget.getInsertHandle(), fragments, statistics, refreshTarget.getSourceTableHandles());
        } else if (target instanceof DeleteTarget) {
            metadata.finishDelete(session, ((DeleteTarget) target).getHandleOrElseThrow(), fragments);
            return Optional.empty();
        } else if (target instanceof UpdateTarget) {
            metadata.finishUpdate(session, ((UpdateTarget) target).getHandleOrElseThrow(), fragments);
            return Optional.empty();
        } else if (target instanceof TableExecuteTarget) {
            TableExecuteHandle tableExecuteHandle = ((TableExecuteTarget) target).getExecuteHandle();
            metadata.finishTableExecute(session, tableExecuteHandle, fragments, tableExecuteContext.getSplitsInfo());
            return Optional.empty();
        } else {
            throw new AssertionError("Unhandled target type: " + target.getClass().getName());
        }
    };
}

Example 2 with UpdateTarget

use of io.trino.sql.planner.plan.TableWriterNode.UpdateTarget 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);
}