diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java index 0d7c27769823b..aaeec08ed1675 100644 --- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java +++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java @@ -21,6 +21,7 @@ import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.analyzer.AccessControlInfo; import com.facebook.presto.spi.analyzer.AccessControlInfoForTable; @@ -32,12 +33,14 @@ import com.facebook.presto.spi.function.FunctionKind; import com.facebook.presto.spi.function.table.Argument; import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; +import com.facebook.presto.spi.plan.PartitioningHandle; import com.facebook.presto.spi.security.AccessControl; import com.facebook.presto.spi.security.AccessControlContext; import com.facebook.presto.spi.security.AllowAllAccessControl; import com.facebook.presto.spi.security.Identity; import com.facebook.presto.sql.tree.ExistsPredicate; import com.facebook.presto.sql.tree.Expression; +import com.facebook.presto.sql.tree.FieldReference; import com.facebook.presto.sql.tree.FunctionCall; import com.facebook.presto.sql.tree.GroupingOperation; import com.facebook.presto.sql.tree.Identifier; @@ -187,6 +190,7 @@ public class Analysis private Optional analyzeTarget = Optional.empty(); private Optional> updatedColumns = Optional.empty(); + private Optional mergeAnalysis = Optional.empty(); // for describe input and describe output private final boolean isDescribe; @@ -221,6 +225,9 @@ public class Analysis private final Set> aliasedRelations = new LinkedHashSet<>(); private final Set> polymorphicTableFunctions = new LinkedHashSet<>(); + // Row id field used for MERGE INTO command. + private final Map, FieldReference> rowIdField = new LinkedHashMap<>(); + public Analysis(@Nullable Statement root, Map, Expression> parameters, boolean isDescribe) { this.root = root; @@ -433,6 +440,16 @@ public Expression getJoinCriteria(Join join) return joins.get(NodeRef.of(join)); } + public void setRowIdField(Table table, FieldReference field) + { + rowIdField.put(NodeRef.of(table), field); + } + + public FieldReference getRowIdField(Table table) + { + return rowIdField.get(NodeRef.of(table)); + } + public void recordSubqueries(Node node, ExpressionAnalysis expressionAnalysis) { NodeRef key = NodeRef.of(node); @@ -726,6 +743,16 @@ public Optional> getUpdatedColumns() return updatedColumns; } + public Optional getMergeAnalysis() + { + return mergeAnalysis; + } + + public void setMergeAnalysis(MergeAnalysis mergeAnalysis) + { + this.mergeAnalysis = Optional.of(mergeAnalysis); + } + public void setRefreshMaterializedViewAnalysis(RefreshMaterializedViewAnalysis refreshMaterializedViewAnalysis) { this.refreshMaterializedViewAnalysis = Optional.of(refreshMaterializedViewAnalysis); @@ -1694,4 +1721,108 @@ public ConnectorTransactionHandle getTransactionHandle() return transactionHandle; } } + + public static class MergeAnalysis + { + private final Table targetTable; + private final List targetColumnsMetadata; + private final List targetColumnHandles; + private final List targetRedistributionColumnHandles; + private final List> mergeCaseColumnHandles; + private final Set nonNullableColumnHandles; + private final Map columnHandleFieldNumbers; + private final List insertPartitioningArgumentIndexes; + private final Optional insertLayout; + private final Optional updateLayout; + private final Scope targetTableScope; + private final Scope joinScope; + + public MergeAnalysis( + Table targetTable, + List targetColumnsMetadata, + List targetColumnHandles, + List targetRedistributionColumnHandles, + List> mergeCaseColumnHandles, + Set nonNullableTargetColumnHandles, + Map targetColumnHandleFieldNumbers, + List insertPartitioningArgumentIndexes, + Optional insertLayout, + Optional updateLayout, + Scope targetTableScope, + Scope joinScope) + { + this.targetTable = requireNonNull(targetTable, "targetTable is null"); + this.targetColumnsMetadata = requireNonNull(targetColumnsMetadata, "targetColumnsMetadata is null"); + this.targetColumnHandles = requireNonNull(targetColumnHandles, "targetColumnHandles is null"); + this.targetRedistributionColumnHandles = requireNonNull(targetRedistributionColumnHandles, "targetRedistributionColumnHandles is null"); + this.mergeCaseColumnHandles = requireNonNull(mergeCaseColumnHandles, "mergeCaseColumnHandles is null"); + this.nonNullableColumnHandles = requireNonNull(nonNullableTargetColumnHandles, "nonNullableTargetColumnHandles is null"); + this.columnHandleFieldNumbers = requireNonNull(targetColumnHandleFieldNumbers, "targetColumnHandleFieldNumbers is null"); + this.insertLayout = requireNonNull(insertLayout, "insertLayout is null"); + this.updateLayout = requireNonNull(updateLayout, "updateLayout is null"); + this.insertPartitioningArgumentIndexes = (requireNonNull(insertPartitioningArgumentIndexes, "insertPartitioningArgumentIndexes is null")); + this.targetTableScope = requireNonNull(targetTableScope, "targetTableScope is null"); + this.joinScope = requireNonNull(joinScope, "joinScope is null"); + } + + public Table getTargetTable() + { + return targetTable; + } + + public List getTargetColumnsMetadata() + { + return targetColumnsMetadata; + } + + public List getTargetColumnHandles() + { + return targetColumnHandles; + } + + public List getTargetRedistributionColumnHandles() + { + return targetRedistributionColumnHandles; + } + + public List> getMergeCaseColumnHandles() + { + return mergeCaseColumnHandles; + } + + public Set getNonNullableColumnHandles() + { + return nonNullableColumnHandles; + } + + public Map getColumnHandleFieldNumbers() + { + return columnHandleFieldNumbers; + } + + public List getInsertPartitioningArgumentIndexes() + { + return insertPartitioningArgumentIndexes; + } + + public Optional getInsertLayout() + { + return insertLayout; + } + + public Optional getUpdateLayout() + { + return updateLayout; + } + + public Scope getJoinScope() + { + return joinScope; + } + + public Scope getTargetTableScope() + { + return targetTableScope; + } + } } diff --git a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java index b125957fe7898..584e137c1c1db 100644 --- a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java +++ b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java @@ -26,6 +26,7 @@ import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import java.util.List; +import java.util.Optional; import java.util.function.ToIntFunction; import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; @@ -43,10 +44,10 @@ public BlackHoleNodePartitioningProvider(NodeManager nodeManager) } @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { // create one bucket per node - return createBucketNodeMap(nodeManager.getRequiredWorkerNodes().size()); + return Optional.of(createBucketNodeMap(nodeManager.getRequiredWorkerNodes().size())); } @Override diff --git a/presto-common/src/main/java/com/facebook/presto/common/Page.java b/presto-common/src/main/java/com/facebook/presto/common/Page.java index 8eb4fe7550cc5..2fac6ad5d5110 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/Page.java +++ b/presto-common/src/main/java/com/facebook/presto/common/Page.java @@ -357,6 +357,17 @@ public Page getLoadedPage(int... channels) return wrapBlocksWithoutCopy(positionCount, blocks); } + public Page getColumns(int... columns) + { + requireNonNull(columns, "columns is null"); + + Block[] blocks = new Block[columns.length]; + for (int i = 0; i < columns.length; i++) { + blocks[i] = this.blocks[columns[i]]; + } + return wrapBlocksWithoutCopy(positionCount, blocks); + } + @Override public String toString() { diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/Block.java b/presto-common/src/main/java/com/facebook/presto/common/block/Block.java index 6b5379ba06ebc..4428e0fbe900b 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/Block.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/Block.java @@ -400,4 +400,20 @@ default long toLong(int position) { throw new UnsupportedOperationException(getClass().getName()); } + + /** + * Returns the underlying value block underlying this block. + */ + default Block getUnderlyingValueBlock() + { + return this; + } + + /** + * Returns the position in the underlying value block corresponding to the specified position in this block. + */ + default int getUnderlyingValuePosition(int position) + { + return position; + } } diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java b/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java index f2d9b3505d4fb..c7eeed336c1d2 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java @@ -307,4 +307,27 @@ static Block[] ensureBlocksAreLoaded(Block[] blocks) // No newly loaded blocks return blocks; } + + static boolean[] copyIsNullAndAppendNull(@Nullable boolean[] isNull, int offsetBase, int positionCount) + { + int desiredLength = offsetBase + positionCount + 1; + boolean[] newIsNull = new boolean[desiredLength]; + if (isNull != null) { + checkArrayRange(isNull, offsetBase, positionCount); + System.arraycopy(isNull, 0, newIsNull, 0, desiredLength - 1); + } + // mark the last element to append null + newIsNull[desiredLength - 1] = true; + return newIsNull; + } + + static int[] copyOffsetsAndAppendNull(int[] offsets, int offsetBase, int positionCount) + { + int desiredLength = offsetBase + positionCount + 2; + checkArrayRange(offsets, offsetBase, positionCount + 1); + int[] newOffsets = Arrays.copyOf(offsets, desiredLength); + // Null element does not move the offset forward + newOffsets[desiredLength - 1] = newOffsets[desiredLength - 2]; + return newOffsets; + } } diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/DictionaryBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/DictionaryBlock.java index 9587f6c8bfc6d..0d0e823535fe8 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/DictionaryBlock.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/DictionaryBlock.java @@ -508,6 +508,45 @@ public Block getLoadedBlock() return new DictionaryBlock(idsOffset, getPositionCount(), loadedDictionary, ids, false, randomDictionaryId()); } + @Override + public Block getUnderlyingValueBlock() + { + return dictionary.getUnderlyingValueBlock(); + } + + @Override + public int getUnderlyingValuePosition(int position) + { + return dictionary.getUnderlyingValuePosition(getId(position)); + } + + public Block createProjection(Block newDictionary) + { + if (newDictionary.getPositionCount() != dictionary.getPositionCount()) { + throw new IllegalArgumentException("newDictionary must have the same position count"); + } + + // if the new dictionary is lazy be careful to not materialize it + if (newDictionary instanceof LazyBlock) { + return new LazyBlock(positionCount, (block) -> { + Block newDictionaryBlock = newDictionary.getBlock(0); + Block newBlock = createProjection(newDictionaryBlock); + block.setBlock(newBlock); + }); + } + if (newDictionary instanceof RunLengthEncodedBlock) { + RunLengthEncodedBlock rle = (RunLengthEncodedBlock) newDictionary; + return new RunLengthEncodedBlock(rle.getValue(), positionCount); + } + + // unwrap dictionary in dictionary + int[] newIds = new int[positionCount]; + for (int position = 0; position < positionCount; position++) { + newIds[position] = newDictionary.getUnderlyingValuePosition(getIdUnchecked(position)); + } + return new DictionaryBlock(0, positionCount, newDictionary.getUnderlyingValueBlock(), newIds, false, randomDictionaryId()); + } + public Block getDictionary() { return dictionary; @@ -533,6 +572,11 @@ public int getId(int position) return ids[position + idsOffset]; } + private int getIdUnchecked(int position) + { + return ids[position + idsOffset]; + } + public DictionaryId getDictionarySourceId() { return dictionarySourceId; diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/LazyBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/LazyBlock.java index 079b29c19ac5f..e3a4903842e43 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/LazyBlock.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/LazyBlock.java @@ -389,6 +389,18 @@ public boolean isNullUnchecked(int internalPosition) return block.isNull(internalPosition); } + @Override + public Block getUnderlyingValueBlock() + { + return block.getUnderlyingValueBlock(); + } + + @Override + public int getUnderlyingValuePosition(int position) + { + return block.getUnderlyingValuePosition(position); + } + @Override public Block appendNull() { diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/RowBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/RowBlock.java index 8fde479c047ca..42ce04160cb20 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/RowBlock.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/RowBlock.java @@ -17,9 +17,11 @@ import org.openjdk.jol.info.ClassLayout; import java.util.Arrays; +import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.function.ObjLongConsumer; +import java.util.stream.Collectors; import static com.facebook.presto.common.block.BlockUtil.ensureBlocksAreLoaded; import static io.airlift.slice.SizeOf.sizeOf; @@ -248,6 +250,39 @@ public void retainedBytesForEachPart(ObjLongConsumer consumer) consumer.accept(this, INSTANCE_SIZE); } + /** + * Returns the row fields from the specified block. The block maybe a LazyBlock, RunLengthEncodedBlock, or + * DictionaryBlock, but the underlying block must be a RowBlock. The returned field blocks will be the same + * length as the specified block, which means they are not null suppressed. + */ + public static List getRowFieldsFromBlock(Block block) + { + // if the block is lazy, be careful to not materialize the nested blocks + if (block instanceof LazyBlock) { + LazyBlock lazyBlock = (LazyBlock) block; + block = lazyBlock.getBlock(0); + } + + if (block instanceof RunLengthEncodedBlock) { + RunLengthEncodedBlock runLengthEncodedBlock = (RunLengthEncodedBlock) block; + RowBlock rowBlock = (RowBlock) runLengthEncodedBlock.getValue(); + return Arrays.stream(rowBlock.fieldBlocks) + .map(fieldBlock -> new RunLengthEncodedBlock(fieldBlock, runLengthEncodedBlock.getPositionCount())) + .collect(Collectors.toList()); + } + if (block instanceof DictionaryBlock) { + DictionaryBlock dictionaryBlock = (DictionaryBlock) block; + RowBlock rowBlock = (RowBlock) dictionaryBlock.getDictionary(); + return Arrays.stream(rowBlock.fieldBlocks) + .map(dictionaryBlock::createProjection) + .collect(Collectors.toList()); + } + if (block instanceof RowBlock) { + return Arrays.asList(((RowBlock) block).fieldBlocks); + } + throw new IllegalArgumentException("Unexpected block type: " + block.getClass().getSimpleName()); + } + @Override public String toString() { diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/RunLengthEncodedBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/RunLengthEncodedBlock.java index 28a553ca9184b..0dc05cf359f8a 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/block/RunLengthEncodedBlock.java +++ b/presto-common/src/main/java/com/facebook/presto/common/block/RunLengthEncodedBlock.java @@ -428,6 +428,18 @@ public Block appendNull() } } + @Override + public Block getUnderlyingValueBlock() + { + return value.getUnderlyingValueBlock(); + } + + @Override + public int getUnderlyingValuePosition(int position) + { + return value.getUnderlyingValuePosition(0); + } + @Override public boolean equals(Object obj) { diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/AbstractType.java b/presto-common/src/main/java/com/facebook/presto/common/type/AbstractType.java index f13fe5552d852..e52c73981fa2a 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/AbstractType.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/AbstractType.java @@ -92,6 +92,12 @@ public void writeBoolean(BlockBuilder blockBuilder, boolean value) throw new UnsupportedOperationException(getClass().getName()); } + @Override + public byte getByte(Block block, int position) + { + throw new UnsupportedOperationException(getClass().getName()); + } + @Override public long getLong(Block block, int position) { diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/FunctionType.java b/presto-common/src/main/java/com/facebook/presto/common/type/FunctionType.java index 611cbde38db63..7da315cd4b6ca 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/FunctionType.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/FunctionType.java @@ -144,6 +144,12 @@ public void writeBoolean(BlockBuilder blockBuilder, boolean value) throw new UnsupportedOperationException(getClass().getName()); } + @Override + public byte getByte(Block block, int position) + { + throw new UnsupportedOperationException(getClass().getName()); + } + @Override public long getLong(Block block, int position) { diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/TinyintType.java b/presto-common/src/main/java/com/facebook/presto/common/type/TinyintType.java index b151f9291e61f..305ba7c71c256 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/TinyintType.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/TinyintType.java @@ -17,6 +17,7 @@ import com.facebook.presto.common.block.Block; import com.facebook.presto.common.block.BlockBuilder; import com.facebook.presto.common.block.BlockBuilderStatus; +import com.facebook.presto.common.block.ByteArrayBlock; import com.facebook.presto.common.block.ByteArrayBlockBuilder; import com.facebook.presto.common.block.PageBuilderStatus; import com.facebook.presto.common.block.UncheckedBlock; @@ -133,6 +134,17 @@ public long getLong(Block block, int position) return (long) block.getByte(position); } + @Override + public byte getByte(Block block, int position) + { + return readByte((ByteArrayBlock) block.getUnderlyingValueBlock(), block.getUnderlyingValuePosition(position)); + } + + private static byte readByte(ByteArrayBlock block, int position) + { + return block.getByte(position); + } + @Override public long getLongUnchecked(UncheckedBlock block, int internalPosition) { diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/Type.java b/presto-common/src/main/java/com/facebook/presto/common/type/Type.java index 251b3b11c055d..e4e2c24a5ca11 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/Type.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/Type.java @@ -98,6 +98,11 @@ default boolean equalValuesAreIdentical() */ boolean getBooleanUnchecked(UncheckedBlock block, int internalPosition); + /** + * Gets the value at the {@code block} {@code position} as a byte. + */ + byte getByte(Block block, int position); + /** * Gets the value at the {@code block} {@code position} as a long. */ diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/TypeWithName.java b/presto-common/src/main/java/com/facebook/presto/common/type/TypeWithName.java index 06d718115d20f..903b31d3e8478 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/TypeWithName.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/TypeWithName.java @@ -145,6 +145,12 @@ public boolean getBooleanUnchecked(UncheckedBlock block, int internalPosition) return type.getBooleanUnchecked(block, internalPosition); } + @Override + public byte getByte(Block block, int position) + { + return type.getByte(block, position); + } + @Override public long getLong(Block block, int position) { diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java index d98ca4ac37323..77d47b2db800d 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java @@ -27,6 +27,7 @@ import com.facebook.presto.spi.schedule.NodeSelectionStrategy; import java.util.List; +import java.util.Optional; import java.util.function.ToIntFunction; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -64,7 +65,7 @@ public BucketFunction getBucketFunction( } @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { HivePartitioningHandle handle = (HivePartitioningHandle) partitioningHandle; NodeSelectionStrategy nodeSelectionStrategy = getNodeSelectionStrategy(session); @@ -72,9 +73,9 @@ public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transa switch (nodeSelectionStrategy) { case HARD_AFFINITY: case SOFT_AFFINITY: - return createBucketNodeMap(Stream.generate(() -> sortedNodes).flatMap(List::stream).limit(bucketCount).collect(toImmutableList()), nodeSelectionStrategy); + return Optional.of(createBucketNodeMap(Stream.generate(() -> sortedNodes).flatMap(List::stream).limit(bucketCount).collect(toImmutableList()), nodeSelectionStrategy)); case NO_PREFERENCE: - return createBucketNodeMap(bucketCount); + return Optional.of(createBucketNodeMap(bucketCount)); default: throw new PrestoException(NODE_SELECTION_NOT_SUPPORTED, format("Unsupported node selection strategy %s", nodeSelectionStrategy)); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/ExecutionWriterTarget.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/ExecutionWriterTarget.java index 2704ab5ab8960..c3a47c96bb46b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/ExecutionWriterTarget.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/ExecutionWriterTarget.java @@ -35,7 +35,8 @@ @JsonSubTypes.Type(value = ExecutionWriterTarget.InsertHandle.class, name = "InsertHandle"), @JsonSubTypes.Type(value = ExecutionWriterTarget.DeleteHandle.class, name = "DeleteHandle"), @JsonSubTypes.Type(value = ExecutionWriterTarget.RefreshMaterializedViewHandle.class, name = "RefreshMaterializedViewHandle"), - @JsonSubTypes.Type(value = ExecutionWriterTarget.UpdateHandle.class, name = "UpdateHandle")}) + @JsonSubTypes.Type(value = ExecutionWriterTarget.UpdateHandle.class, name = "UpdateHandle"), + @JsonSubTypes.Type(value = ExecutionWriterTarget.MergeHandle.class, name = "MergeHandle")}) @SuppressWarnings({"EmptyClass", "ClassMayBeInterface"}) public abstract class ExecutionWriterTarget { @@ -228,4 +229,48 @@ public String toString() return handle.toString(); } } + + public static class MergeHandle + extends ExecutionWriterTarget + { + private final com.facebook.presto.spi.MergeHandle handle; + // TODO #20578: Uncomment if finally it is necessary. +// private final SchemaTableName schemaTableName; +// private final ConnectorMergeTableHandle connectorMergeTableHandle; + + @JsonCreator + public MergeHandle( + @JsonProperty("handle") com.facebook.presto.spi.MergeHandle handle) +// @JsonProperty("schemaTableName") SchemaTableName schemaTableName, +// @JsonProperty("connectorMergeTableHandle") ConnectorMergeTableHandle connectorMergeTableHandle) + { + this.handle = requireNonNull(handle, "tableHandle is null"); +// this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null"); +// this.connectorMergeTableHandle = requireNonNull(connectorMergeTableHandle, "connectorMergeTableHandle is null"); + } + + @JsonProperty + public com.facebook.presto.spi.MergeHandle getHandle() + { + return handle; + } + +// @JsonProperty +// public SchemaTableName getSchemaTableName() +// { +// return schemaTableName; +// } + +// @JsonProperty +// public ConnectorMergeTableHandle getConnectorMergeTableHandle() +// { +// return connectorMergeTableHandle; +// } + + @Override + public String toString() + { + return handle.toString(); + } + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/TableWriteInfo.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/TableWriteInfo.java index 2d2fa96b6882e..beabe411d1889 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/TableWriteInfo.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/TableWriteInfo.java @@ -20,6 +20,7 @@ import com.facebook.presto.Session; import com.facebook.presto.metadata.AnalyzeTableHandle; import com.facebook.presto.metadata.Metadata; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.plan.PlanNode; import com.facebook.presto.spi.plan.TableFinishNode; import com.facebook.presto.spi.plan.TableWriterNode; @@ -101,6 +102,12 @@ private static Optional createWriterTarget(Optional mergeHandle = mergeTarget.getMergeHandle(); + return Optional.of(new ExecutionWriterTarget.MergeHandle(mergeHandle.orElseThrow( + () -> new VerifyException("mergeHandle is absent: " + target.getClass().getSimpleName())))); + } throw new IllegalArgumentException("Unhandled target type: " + target.getClass().getSimpleName()); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java index 45c1cc7af319b..f43f16a36378a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java @@ -26,6 +26,7 @@ import com.facebook.presto.spi.ConnectorTableMetadata; import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.MaterializedViewDefinition; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.SystemTable; import com.facebook.presto.spi.TableHandle; @@ -34,6 +35,7 @@ import com.facebook.presto.spi.connector.ConnectorCapabilities; import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; @@ -412,6 +414,36 @@ public void finishUpdate(Session session, TableHandle tableHandle, Collection getMergeUpdateLayout(Session session, TableHandle tableHandle) + { + return delegate.getMergeUpdateLayout(session, tableHandle); + } + + @Override + public MergeHandle beginMerge(Session session, TableHandle tableHandle) + { + return delegate.beginMerge(session, tableHandle); + } + + @Override + public void finishMerge(Session session, MergeHandle tableHandle, Collection fragments, Collection computedStatistics) + { + delegate.finishMerge(session, tableHandle, fragments, computedStatistics); + } + @Override public Optional getCatalogHandle(Session session, String catalogName) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleJsonModule.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleJsonModule.java index 952fc38be9426..691ef0146c4c0 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleJsonModule.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleJsonModule.java @@ -33,6 +33,7 @@ public void configure(Binder binder) jsonBinder(binder).addModuleBinding().to(OutputTableHandleJacksonModule.class); jsonBinder(binder).addModuleBinding().to(InsertTableHandleJacksonModule.class); jsonBinder(binder).addModuleBinding().to(DeleteTableHandleJacksonModule.class); + jsonBinder(binder).addModuleBinding().to(MergeHandleJacksonModule.class); jsonBinder(binder).addModuleBinding().to(IndexHandleJacksonModule.class); jsonBinder(binder).addModuleBinding().to(TransactionHandleJacksonModule.class); jsonBinder(binder).addModuleBinding().to(PartitioningHandleJacksonModule.class); diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleResolver.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleResolver.java index 30630f0a7bff7..c6d3bb6015dd0 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleResolver.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/HandleResolver.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.ConnectorHandleResolver; import com.facebook.presto.spi.ConnectorIndexHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorSplit; import com.facebook.presto.spi.ConnectorTableHandle; @@ -134,6 +135,11 @@ public String getId(FunctionHandle functionHandle) return getFunctionNamespaceId(functionHandle, MaterializedFunctionHandleResolver::getFunctionHandleClass); } + public String getId(ConnectorMergeTableHandle mergeHandle) + { + return getId(mergeHandle, MaterializedHandleResolver::getMergeHandleClass); + } + public Class getTableHandleClass(String id) { return resolverFor(id).getTableHandleClass().orElseThrow(() -> new IllegalArgumentException("No resolver for " + id)); @@ -174,6 +180,11 @@ public Class getDeleteTableHandleClass(Str return resolverFor(id).getDeleteTableHandleClass().orElseThrow(() -> new IllegalArgumentException("No resolver for " + id)); } + public Class getMergeHandleClass(String id) + { + return resolverFor(id).getMergeHandleClass().orElseThrow(() -> new IllegalArgumentException("No resolver for " + id)); + } + public Class getPartitioningHandleClass(String id) { return resolverFor(id).getPartitioningHandleClass().orElseThrow(() -> new IllegalArgumentException("No resolver for " + id)); @@ -241,6 +252,7 @@ private static class MaterializedHandleResolver private final Optional> outputTableHandle; private final Optional> insertTableHandle; private final Optional> deleteTableHandle; + private final Optional> mergeTableHandle; private final Optional> partitioningHandle; private final Optional> transactionHandle; @@ -254,6 +266,7 @@ public MaterializedHandleResolver(ConnectorHandleResolver resolver) outputTableHandle = getHandleClass(resolver::getOutputTableHandleClass); insertTableHandle = getHandleClass(resolver::getInsertTableHandleClass); deleteTableHandle = getHandleClass(resolver::getDeleteTableHandleClass); + mergeTableHandle = getHandleClass(resolver::getMergeTableHandleClass); partitioningHandle = getHandleClass(resolver::getPartitioningHandleClass); transactionHandle = getHandleClass(resolver::getTransactionHandleClass); } @@ -308,6 +321,11 @@ public Optional> getDeleteTableHandl return deleteTableHandle; } + public Optional> getMergeHandleClass() + { + return mergeTableHandle; + } + public Optional> getPartitioningHandleClass() { return partitioningHandle; @@ -336,6 +354,7 @@ public boolean equals(Object o) Objects.equals(outputTableHandle, that.outputTableHandle) && Objects.equals(insertTableHandle, that.insertTableHandle) && Objects.equals(deleteTableHandle, that.deleteTableHandle) && + Objects.equals(mergeTableHandle, that.mergeTableHandle) && Objects.equals(partitioningHandle, that.partitioningHandle) && Objects.equals(transactionHandle, that.transactionHandle); } @@ -343,7 +362,7 @@ public boolean equals(Object o) @Override public int hashCode() { - return Objects.hash(tableHandle, layoutHandle, columnHandle, split, indexHandle, outputTableHandle, insertTableHandle, deleteTableHandle, partitioningHandle, transactionHandle); + return Objects.hash(tableHandle, layoutHandle, columnHandle, split, indexHandle, outputTableHandle, insertTableHandle, deleteTableHandle, mergeTableHandle, partitioningHandle, transactionHandle); } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MergeHandleJacksonModule.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MergeHandleJacksonModule.java new file mode 100644 index 0000000000000..1669b22d5e7f5 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MergeHandleJacksonModule.java @@ -0,0 +1,30 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.metadata; + +import com.facebook.presto.spi.ConnectorMergeTableHandle; + +import javax.inject.Inject; + +public class MergeHandleJacksonModule + extends AbstractTypedJacksonModule +{ + @Inject + public MergeHandleJacksonModule(HandleResolver handleResolver) + { + super(ConnectorMergeTableHandle.class, + handleResolver::getId, + handleResolver::getMergeHandleClass); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java index 050702ff6ad38..3831cf44e6d54 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java @@ -28,6 +28,7 @@ import com.facebook.presto.spi.ConnectorTableMetadata; import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.MaterializedViewDefinition; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.SystemTable; @@ -41,6 +42,7 @@ import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; @@ -351,6 +353,34 @@ public interface Metadata */ void finishUpdate(Session session, TableHandle tableHandle, Collection fragments); + /** + * Return the row update paradigm supported by the connector on the table or throw + * an exception if row change is not supported. + */ + RowChangeParadigm getRowChangeParadigm(Session session, TableHandle tableHandle); + + /** + * Get the column handle that will generate row IDs for the merge operation. + * These IDs will be passed to the {@code storeMergedRows()} method of the + * {@link com.facebook.presto.spi.ConnectorMergeSink} that created them. + */ + ColumnHandle getMergeRowIdColumnHandle(Session session, TableHandle tableHandle); + + /** + * Get the physical layout for updated rows of a MERGE operation. + */ + Optional getMergeUpdateLayout(Session session, TableHandle tableHandle); + + /** + * Begin merge query + */ + MergeHandle beginMerge(Session session, TableHandle tableHandle); + + /** + * Finish merge query + */ + void finishMerge(Session session, MergeHandle tableHandle, Collection fragments, Collection computedStatistics); + /** * Returns a connector id for the specified catalog name. */ diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java index a1be6217dbb60..c6bcfb31cac88 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java @@ -32,6 +32,7 @@ import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorResolvedIndex; import com.facebook.presto.spi.ConnectorSession; @@ -43,6 +44,7 @@ import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.MaterializedViewDefinition; import com.facebook.presto.spi.MaterializedViewStatus; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.QueryId; @@ -61,6 +63,7 @@ import com.facebook.presto.spi.connector.ConnectorPartitioningMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; @@ -921,6 +924,26 @@ public Optional getUpdateRowIdColumn(Session session, TableHandle return metadata.getUpdateRowIdColumn(session.toConnectorSession(connectorId), tableHandle.getConnectorHandle(), updatedColumns); } + @Override + public ColumnHandle getMergeRowIdColumnHandle(Session session, TableHandle tableHandle) + { + ConnectorId connectorId = tableHandle.getConnectorId(); + ConnectorMetadata metadata = getMetadata(session, connectorId); + return metadata.getMergeRowIdColumnHandle(session.toConnectorSession(connectorId), tableHandle.getConnectorHandle()); + } + + @Override + public Optional getMergeUpdateLayout(Session session, TableHandle tableHandle) + { + ConnectorId connectorId = tableHandle.getConnectorId(); + CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, connectorId); + ConnectorMetadata metadata = catalogMetadata.getMetadata(); + ConnectorTransactionHandle transactionHandle = catalogMetadata.getTransactionHandleFor(connectorId); + + return metadata.getMergeUpdateLayout(session.toConnectorSession(connectorId), tableHandle.getConnectorHandle()) + .map(partitioning -> new PartitioningHandle(Optional.of(connectorId), Optional.of(transactionHandle), partitioning)); + } + @Override public boolean supportsMetadataDelete(Session session, TableHandle tableHandle) { @@ -977,6 +1000,35 @@ public void finishUpdate(Session session, TableHandle tableHandle, Collection fragments, + Collection computedStatistics) + { + ConnectorId connectorId = mergeHandle.getTableHandle().getConnectorId(); + ConnectorMetadata metadata = getMetadata(session, connectorId); + metadata.finishMerge(session.toConnectorSession(connectorId), mergeHandle.getConnectorMergeTableHandle(), fragments, computedStatistics); + } + @Override public Optional getCatalogHandle(Session session, String catalogName) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/AbstractRowChangeOperator.java b/presto-main-base/src/main/java/com/facebook/presto/operator/AbstractRowChangeOperator.java index 1fd421543539c..f26ba51630fba 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/AbstractRowChangeOperator.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/AbstractRowChangeOperator.java @@ -48,7 +48,7 @@ protected enum State protected State state = State.RUNNING; protected long rowCount; private boolean closed; - private ListenableFuture> finishFuture; + protected ListenableFuture> finishFuture; private Supplier> pageSource = Optional::empty; private final JsonCodec tableCommitContextCodec; @@ -158,6 +158,7 @@ public void close() } else { pageSource.get().ifPresent(UpdatablePageSource::abort); + abort(); } } } @@ -173,4 +174,6 @@ protected UpdatablePageSource pageSource() // empty source can occur if the source operator doesn't output any rows return source.orElseGet(EmptySplitPageSource::new); } + + protected void abort() {} } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java b/presto-main-base/src/main/java/com/facebook/presto/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java new file mode 100644 index 0000000000000..a16cd28d44ecf --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java @@ -0,0 +1,110 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator; + +import com.facebook.presto.common.Page; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.common.block.RunLengthEncodedBlock; + +import java.util.ArrayList; +import java.util.List; + +import static com.facebook.presto.common.Utils.nativeValueToBlock; +import static com.facebook.presto.common.block.RowBlock.getRowFieldsFromBlock; +import static com.facebook.presto.common.type.TinyintType.TINYINT; +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +/** + * The transformPage() method in this class does two things: + *
    + *
  • Transform the input page into an "update" page format
    • + *
  • Removes all rows whose operation number is DEFAULT_CASE_OPERATION_NUMBER
    • + *
+ */ +public class ChangeOnlyUpdatedColumnsMergeProcessor + implements MergeRowChangeProcessor +{ + private static final Block INSERT_FROM_UPDATE_BLOCK = nativeValueToBlock(TINYINT, 0L); + + private final int rowIdChannel; + private final int mergeRowChannel; + private final List dataColumnChannels; + private final int writeRedistributionColumnCount; + + public ChangeOnlyUpdatedColumnsMergeProcessor( + int rowIdChannel, + int mergeRowChannel, + List targetColumnChannels, + List redistributionColumnChannels) + { + this.rowIdChannel = rowIdChannel; + this.mergeRowChannel = mergeRowChannel; + this.dataColumnChannels = requireNonNull(targetColumnChannels, "targetColumnChannels is null"); + this.writeRedistributionColumnCount = redistributionColumnChannels.size(); + } + + @Override + public Page transformPage(Page inputPage) + { + requireNonNull(inputPage, "inputPage is null"); + + int inputChannelCount = inputPage.getChannelCount(); + checkArgument(inputChannelCount >= 2 + writeRedistributionColumnCount, "inputPage channelCount (%s) should be >= 2 + %s", inputChannelCount, writeRedistributionColumnCount); + int positionCount = inputPage.getPositionCount(); + checkArgument(positionCount > 0, "positionCount should be > 0, but is %s", positionCount); + + Block mergeRow = inputPage.getBlock(mergeRowChannel).getLoadedBlock(); + if (mergeRow.mayHaveNull()) { + for (int position = 0; position < positionCount; position++) { + checkArgument(!mergeRow.isNull(position), "The mergeRow may not have null rows"); + } + } + + List fields = getRowFieldsFromBlock(mergeRow); + List builder = new ArrayList<>(dataColumnChannels.size() + 3); + for (int channel : dataColumnChannels) { + builder.add(fields.get(channel)); + } + Block operationChannelBlock = fields.get(fields.size() - 2); + builder.add(operationChannelBlock); + builder.add(inputPage.getBlock(rowIdChannel)); + builder.add(new RunLengthEncodedBlock(INSERT_FROM_UPDATE_BLOCK, positionCount)); + + Page result = new Page(builder.toArray(new Block[0])); + + int defaultCaseCount = 0; + for (int position = 0; position < positionCount; position++) { + if (TINYINT.getByte(operationChannelBlock, position) == DEFAULT_CASE_OPERATION_NUMBER) { + defaultCaseCount++; + } + } + if (defaultCaseCount == 0) { + return result; + } + + int usedCases = 0; + int[] positions = new int[positionCount - defaultCaseCount]; + for (int position = 0; position < positionCount; position++) { + if (TINYINT.getByte(operationChannelBlock, position) != DEFAULT_CASE_OPERATION_NUMBER) { + positions[usedCases] = position; + usedCases++; + } + } + + checkArgument(usedCases + defaultCaseCount == positionCount, "usedCases (%s) + defaultCaseCount (%s) != positionCount (%s)", usedCases, defaultCaseCount, positionCount); + + return result.getPositions(positions, 0, usedCases); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/DeleteAndInsertMergeProcessor.java b/presto-main-base/src/main/java/com/facebook/presto/operator/DeleteAndInsertMergeProcessor.java new file mode 100644 index 0000000000000..6fb9240c1f4ea --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/DeleteAndInsertMergeProcessor.java @@ -0,0 +1,205 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator; + +import com.facebook.presto.common.Page; +import com.facebook.presto.common.PageBuilder; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.block.ColumnarRow; +import com.facebook.presto.common.type.Type; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +import static com.facebook.presto.common.block.ColumnarRow.toColumnarRow; +import static com.facebook.presto.common.type.TinyintType.TINYINT; +import static com.facebook.presto.spi.ConnectorMergeSink.DELETE_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.INSERT_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.UPDATE_OPERATION_NUMBER; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; +import static java.util.Objects.requireNonNull; + +public class DeleteAndInsertMergeProcessor + implements MergeRowChangeProcessor +{ + private final List targetColumnTypes; + private final Type rowIdType; + private final int rowIdChannel; + private final int mergeRowChannel; + private final List targetColumnChannels; + private final int redistributionColumnCount; + private final List redistributionChannelNumbers; + + public DeleteAndInsertMergeProcessor( + List targetColumnTypes, + Type rowIdType, + int rowIdChannel, + int mergeRowChannel, + List redistributionChannelNumbers, + List targetColumnChannels) + { + this.targetColumnTypes = requireNonNull(targetColumnTypes, "targetColumnTypes is null"); + this.rowIdType = requireNonNull(rowIdType, "rowIdType is null"); + this.rowIdChannel = rowIdChannel; + this.mergeRowChannel = mergeRowChannel; + this.redistributionColumnCount = redistributionChannelNumbers.size(); + int redistributionSourceIndex = 0; + this.targetColumnChannels = requireNonNull(targetColumnChannels, "targetColumnChannels is null"); + ImmutableList.Builder redistributionChannelNumbersBuilder = ImmutableList.builder(); + for (int dataColumnChannel : targetColumnChannels) { + if (redistributionChannelNumbers.contains(dataColumnChannel)) { + redistributionChannelNumbersBuilder.add(redistributionSourceIndex); + redistributionSourceIndex++; + } + else { + redistributionChannelNumbersBuilder.add(-1); + } + } + this.redistributionChannelNumbers = redistributionChannelNumbersBuilder.build(); + } + + @JsonProperty + public Type getRowIdType() + { + return rowIdType; + } + + /** + * Transform UPDATE operations into an INSERT and DELETE operation. + * See {@link MergeRowChangeProcessor#transformPage} for details. + * @param inputPage It has 5 channels/blocks:
+ * 1. Unique ID
+ * 2. Merge Row ID (_file:varchar, _pos:bigint, file_record_count:bigint, partition_spec_id:integer, partition_data:varchar)
+ * 3. Merge Row (source table columns, is present, operation, case number)
+ * 4. Merge case number
+ * 5. Is distinct row: it is 1 if no other row has the same unique id and WHEN clause number, 0 otherwise.
+ */ + @Override + public Page transformPage(Page inputPage) + { + requireNonNull(inputPage, "inputPage is null"); + int inputChannelCount = inputPage.getChannelCount(); + checkArgument(inputChannelCount >= 2 + redistributionColumnCount, "inputPage channelCount (%s) should be >= 2 + partition columns size (%s)", inputChannelCount, redistributionColumnCount); + + int originalPositionCount = inputPage.getPositionCount(); + checkArgument(originalPositionCount > 0, "originalPositionCount should be > 0, but is %s", originalPositionCount); + + ColumnarRow mergeRow = toColumnarRow(inputPage.getBlock(mergeRowChannel)); + Block operationChannelBlock = mergeRow.getField(mergeRow.getFieldCount() - 2); + + int updatePositions = 0; + int insertPositions = 0; + int deletePositions = 0; + for (int position = 0; position < originalPositionCount; position++) { + byte operation = TINYINT.getByte(operationChannelBlock, position); + switch (operation) { + case DEFAULT_CASE_OPERATION_NUMBER:/* ignored */ + break; + case INSERT_OPERATION_NUMBER: + insertPositions++; + break; + case DELETE_OPERATION_NUMBER: + deletePositions++; + break; + case UPDATE_OPERATION_NUMBER: + updatePositions++; + break; + default: + throw new IllegalArgumentException("Unknown operator number: " + operation); + } + } + + int totalPositions = insertPositions + deletePositions + (2 * updatePositions); + List pageTypes = ImmutableList.builder() + .addAll(targetColumnTypes) + .add(TINYINT) // Operation: INSERT(1), DELETE(2), UPDATE(3). More info: ConnectorMergeSink + .add(rowIdType) + .add(TINYINT) // Insert from update: it is 1 if the cause of the insert is an UPDATE, 0 otherwise. + .build(); + + PageBuilder pageBuilder = new PageBuilder(totalPositions, pageTypes); + for (int position = 0; position < originalPositionCount; position++) { + byte operation = TINYINT.getByte(operationChannelBlock, position); + if (operation != DEFAULT_CASE_OPERATION_NUMBER) { + // Delete and Update because both create a delete row + if (operation == DELETE_OPERATION_NUMBER || operation == UPDATE_OPERATION_NUMBER) { + addDeleteRow(pageBuilder, inputPage, position); + } + // Insert and update because both create an insert row + if (operation == INSERT_OPERATION_NUMBER || operation == UPDATE_OPERATION_NUMBER) { + addInsertRow(pageBuilder, mergeRow, position, operation == UPDATE_OPERATION_NUMBER); + } + } + } + + Page page = pageBuilder.build(); + verify(page.getPositionCount() == totalPositions, "page positions (%s) is not equal to (%s)", page.getPositionCount(), totalPositions); + return page; + } + + private void addDeleteRow(PageBuilder pageBuilder, Page originalPage, int position) + { + // Copy the write redistribution columns + for (int targetChannel : targetColumnChannels) { + Type columnType = targetColumnTypes.get(targetChannel); + BlockBuilder targetBlock = pageBuilder.getBlockBuilder(targetChannel); + + int redistributionChannelNumber = redistributionChannelNumbers.get(targetChannel); + if (redistributionChannelNumbers.get(targetChannel) >= 0) { + // The value comes from that column of the page + columnType.appendTo(originalPage.getBlock(redistributionChannelNumber), position, targetBlock); + } + else { + // We don't care about the other data columns + targetBlock.appendNull(); + } + } + + // Add the operation column == deleted + TINYINT.writeLong(pageBuilder.getBlockBuilder(targetColumnChannels.size()), DELETE_OPERATION_NUMBER); + + // Copy row ID column + rowIdType.appendTo(originalPage.getBlock(rowIdChannel), position, pageBuilder.getBlockBuilder(targetColumnChannels.size() + 1)); + + // Write 0, meaning this row is not an insert derived from an update + TINYINT.writeLong(pageBuilder.getBlockBuilder(targetColumnChannels.size() + 2), 0); + + pageBuilder.declarePosition(); + } + + private void addInsertRow(PageBuilder pageBuilder, ColumnarRow mergeCaseBlock, int position, boolean causedByUpdate) + { + // Copy the values from the merge block + for (int targetChannel : targetColumnChannels) { + Type columnType = targetColumnTypes.get(targetChannel); + BlockBuilder targetBlock = pageBuilder.getBlockBuilder(targetChannel); + // The value comes from that column of the page + columnType.appendTo(mergeCaseBlock.getField(targetChannel), position, targetBlock); + } + + // Add the operation column == insert + TINYINT.writeLong(pageBuilder.getBlockBuilder(targetColumnChannels.size()), INSERT_OPERATION_NUMBER); + + // Add null row ID column + pageBuilder.getBlockBuilder(targetColumnChannels.size() + 1).appendNull(); + + // Write 1 if this row is an insert derived from an update, 0 otherwise + TINYINT.writeLong(pageBuilder.getBlockBuilder(targetColumnChannels.size() + 2), causedByUpdate ? 1 : 0); + + pageBuilder.declarePosition(); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/Driver.java b/presto-main-base/src/main/java/com/facebook/presto/operator/Driver.java index 442572f7725fb..9cf41c0f563ca 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/Driver.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/Driver.java @@ -81,6 +81,7 @@ public class Driver private final Optional sourceOperator; private final Optional deleteOperator; private final Optional updateOperator; + private final Optional mergeOperator; // This variable acts as a staging area. When new splits (encapsulated in TaskSource) are // provided to a Driver, the Driver will not process them right away. Instead, the splits are @@ -141,6 +142,7 @@ private Driver(DriverContext driverContext, List operators) Optional sourceOperator = Optional.empty(); Optional deleteOperator = Optional.empty(); Optional updateOperator = Optional.empty(); + Optional mergeOperator = Optional.empty(); for (Operator operator : operators) { if (operator instanceof SourceOperator) { checkArgument(!sourceOperator.isPresent(), "There must be at most one SourceOperator"); @@ -154,10 +156,15 @@ else if (operator instanceof UpdateOperator) { checkArgument(!updateOperator.isPresent(), "There must be at most one UpdateOperator"); updateOperator = Optional.of((UpdateOperator) operator); } + else if (operator instanceof MergeWriterOperator) { + checkArgument(!mergeOperator.isPresent(), "There must be at most one MergeWriterOperator"); + mergeOperator = Optional.of((MergeWriterOperator) operator); + } } this.sourceOperator = sourceOperator; this.deleteOperator = deleteOperator; this.updateOperator = updateOperator; + this.mergeOperator = mergeOperator; currentTaskSource = sourceOperator.map(operator -> new TaskSource(operator.getSourceId(), ImmutableSet.of(), false)).orElse(null); // initially the driverBlockedFuture is not blocked (it is completed) @@ -289,6 +296,7 @@ private void processNewSources() Supplier> pageSource = sourceOperator.addSplit(newSplit); deleteOperator.ifPresent(deleteOperator -> deleteOperator.setPageSource(pageSource)); updateOperator.ifPresent(updateOperator -> updateOperator.setPageSource(pageSource)); + mergeOperator.ifPresent(mergeOperator -> mergeOperator.setPageSource(pageSource)); } // set no more splits diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/MergeProcessorOperator.java b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeProcessorOperator.java new file mode 100644 index 0000000000000..2c3537d36212a --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeProcessorOperator.java @@ -0,0 +1,176 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator; + +import com.facebook.presto.common.Page; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.plan.PlanNodeId; +import com.facebook.presto.spi.plan.TableWriterNode.MergeParadigmAndTypes; + +import java.util.List; + +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + +/** + * This operator is used by operations like SQL MERGE. It is used + * for all {@link com.facebook.presto.spi.connector.RowChangeParadigm}s. This operator + * creates the {@link MergeRowChangeProcessor}. + */ +public class MergeProcessorOperator + implements Operator +{ + public static class MergeProcessorOperatorFactory + implements OperatorFactory + { + private final int operatorId; + private final PlanNodeId planNodeId; + private final MergeRowChangeProcessor rowChangeProcessor; + private boolean closed; + + private MergeProcessorOperatorFactory( + int operatorId, + PlanNodeId planNodeId, + MergeRowChangeProcessor rowChangeProcessor) + { + this.operatorId = operatorId; + this.planNodeId = requireNonNull(planNodeId, "planNodeId is null"); + this.rowChangeProcessor = requireNonNull(rowChangeProcessor, "rowChangeProcessor is null"); + } + + public MergeProcessorOperatorFactory( + int operatorId, + PlanNodeId planNodeId, + MergeParadigmAndTypes merge, + int rowIdChannel, + int mergeRowChannel, + List redistributionColumns, + List targetColumnChannels) + { + MergeRowChangeProcessor rowChangeProcessor = createRowChangeProcessor(merge, rowIdChannel, mergeRowChannel, redistributionColumns, targetColumnChannels); + + this.operatorId = operatorId; + this.planNodeId = requireNonNull(planNodeId, "planNodeId is null"); + this.rowChangeProcessor = requireNonNull(rowChangeProcessor, "rowChangeProcessor is null"); + } + + private static MergeRowChangeProcessor createRowChangeProcessor( + MergeParadigmAndTypes merge, + int rowIdChannel, + int mergeRowChannel, + List redistributionColumnChannels, + List targetColumnChannels) + { + switch (merge.getParadigm()) { + case DELETE_ROW_AND_INSERT_ROW: + return new DeleteAndInsertMergeProcessor( + merge.getColumnTypes(), + merge.getRowIdType(), + rowIdChannel, + mergeRowChannel, + redistributionColumnChannels, + targetColumnChannels); + case CHANGE_ONLY_UPDATED_COLUMNS: + return new ChangeOnlyUpdatedColumnsMergeProcessor( + rowIdChannel, + mergeRowChannel, + targetColumnChannels, + redistributionColumnChannels); + default: + throw new PrestoException(NOT_SUPPORTED, "Merge paradigm not supported: " + merge.getParadigm()); + } + } + + @Override + public Operator createOperator(DriverContext driverContext) + { + checkState(!closed, "Factory is already closed"); + OperatorContext context = driverContext.addOperatorContext(operatorId, planNodeId, MergeProcessorOperator.class.getSimpleName()); + return new MergeProcessorOperator(context, rowChangeProcessor); + } + + @Override + public void noMoreOperators() + { + closed = true; + } + + @Override + public OperatorFactory duplicate() + { + return new MergeProcessorOperatorFactory(operatorId, planNodeId, rowChangeProcessor); + } + } + + private final OperatorContext operatorContext; + private final MergeRowChangeProcessor rowChangeProcessor; + + private Page currentPage; + private boolean finishing; + + public MergeProcessorOperator( + OperatorContext operatorContext, + MergeRowChangeProcessor rowChangeProcessor) + { + this.operatorContext = requireNonNull(operatorContext, "operatorContext is null"); + this.rowChangeProcessor = requireNonNull(rowChangeProcessor, "rowChangeProcessor is null"); + } + + @Override + public OperatorContext getOperatorContext() + { + return operatorContext; + } + + @Override + public void finish() + { + finishing = true; + } + + @Override + public boolean isFinished() + { + return finishing && currentPage == null; + } + + @Override + public boolean needsInput() + { + return !finishing && currentPage == null; + } + + @Override + public void addInput(Page page) + { + checkState(!finishing, "Operator is already finishing"); + checkState(currentPage == null, "currentPage must be null to add a new page"); + + currentPage = requireNonNull(page, "page is null"); + } + + @Override + public Page getOutput() + { + if (currentPage == null) { + return null; + } + + Page transformedPage = rowChangeProcessor.transformPage(currentPage); + currentPage = null; + + return transformedPage; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/MergeRowChangeProcessor.java b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeRowChangeProcessor.java new file mode 100644 index 0000000000000..e78bac39413e1 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeRowChangeProcessor.java @@ -0,0 +1,45 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator; + +import com.facebook.presto.common.Page; +import com.facebook.presto.spi.ConnectorMergeSink; + +public interface MergeRowChangeProcessor +{ + int DEFAULT_CASE_OPERATION_NUMBER = -1; + + /** + * Transform a page generated by an SQL MERGE operation into page of data columns and + * operations. The SQL MERGE input page consists of the following: + *
    + *
  • The write redistribution columns, if any
    • + *
  • For partitioned or bucketed tables, a hash value column
    • + *
  • The rowId column for the row from the target table if matched, or null if not matched
    • + *
  • The merge case row block
    • + *
+ * The output page consists of the following: + *
    + *
  • All data columns, in table column order
    • + *
  • {@link ConnectorMergeSink#storeMergedRows The operation block}
    • + *
  • The rowId block
    • + *
  • The last column in the resulting page is 1 if the row is an insert + * derived from an update, and zero otherwise.
    • + *
+ *

+ * The {@link DeleteAndInsertMergeProcessor} implementation will transform each UPDATE + * row into multiple rows: an INSERT row and a DELETE row. + */ + Page transformPage(Page inputPage); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/MergeWriterOperator.java b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeWriterOperator.java new file mode 100644 index 0000000000000..2ba54fcc57938 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/MergeWriterOperator.java @@ -0,0 +1,141 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.presto.Session; +import com.facebook.presto.common.Page; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.spi.ConnectorMergeSink; +import com.facebook.presto.spi.plan.PlanNodeId; +import com.facebook.presto.spi.plan.TableWriterNode.MergeTarget; +import com.facebook.presto.split.PageSinkManager; + +import java.util.stream.IntStream; + +import static com.facebook.airlift.concurrent.MoreFutures.toListenableFuture; +import static com.facebook.presto.common.type.TinyintType.TINYINT; +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + +public class MergeWriterOperator + extends AbstractRowChangeOperator +{ + public static class MergeWriterOperatorFactory + implements OperatorFactory + { + private final int operatorId; + private final PlanNodeId planNodeId; + private final PageSinkManager pageSinkManager; + private final MergeTarget target; + private final Session session; + private final JsonCodec tableCommitContextCodec; + private boolean closed; + + public MergeWriterOperatorFactory( + int operatorId, + PlanNodeId planNodeId, + PageSinkManager pageSinkManager, + MergeTarget target, + Session session, + JsonCodec tableCommitContextCodec) + { + this.operatorId = operatorId; + this.planNodeId = requireNonNull(planNodeId, "planNodeId is null"); + this.pageSinkManager = requireNonNull(pageSinkManager, "pageSinkManager is null"); + this.target = requireNonNull(target, "target is null"); + this.session = requireNonNull(session, "session is null"); + this.tableCommitContextCodec = requireNonNull(tableCommitContextCodec, "tableCommitContextCodec is null"); + } + + @Override + public Operator createOperator(DriverContext driverContext) + { + checkState(!closed, "Factory is already closed"); + OperatorContext context = driverContext.addOperatorContext(operatorId, planNodeId, MergeWriterOperator.class.getSimpleName()); + ConnectorMergeSink mergeSink = pageSinkManager.createMergeSink(session, target.getMergeHandle().get()); + return new MergeWriterOperator(context, mergeSink, tableCommitContextCodec); + } + + @Override + public void noMoreOperators() + { + closed = true; + } + + @Override + public OperatorFactory duplicate() + { + return new MergeWriterOperatorFactory(operatorId, planNodeId, pageSinkManager, target, session, tableCommitContextCodec); + } + } + + private final ConnectorMergeSink mergeSink; + + public MergeWriterOperator(OperatorContext operatorContext, ConnectorMergeSink mergeSink, JsonCodec tableCommitContextCodec) + { + super(operatorContext, tableCommitContextCodec); + this.mergeSink = requireNonNull(mergeSink, "mergeSink is null"); + } + + /** + * @param page It has N + 3 channels/blocks, where N is the number of columns in the source table.
+ * 1: Source table column 1.
+ * 2: Source table column 2.
+ * N: Source table column N.
+ * N + 1: Operation: INSERT(1), DELETE(2), UPDATE(3). More info: {@link ConnectorMergeSink}
+ * N + 2: Merge Row ID (_file:varchar, _pos:bigint, file_record_count:bigint, partition_spec_id:integer, partition_data:varchar).
+ * N + 3: Insert from update: it is 1 if the cause of the insert is an UPDATE, 0 otherwise.
+ */ + @Override + public void addInput(Page page) + { + requireNonNull(page, "page is null"); + checkState(state == State.RUNNING, "Operator is %s", state); + + // Copy all but the last block to a new page. + // The last block exists only to get the rowCount right. + int outputChannelCount = page.getChannelCount() - 1; + + int[] columns = IntStream.range(0, outputChannelCount).toArray(); + Page newPage = page.getColumns(columns); + + // Store the page + mergeSink.storeMergedRows(newPage); + + // Calculate the amount to increment the rowCount + Block insertFromUpdateColumn = page.getBlock(page.getChannelCount() - 1); + long insertsFromUpdates = 0; + int positionCount = page.getPositionCount(); + for (int position = 0; position < positionCount; position++) { + insertsFromUpdates += TINYINT.getByte(insertFromUpdateColumn, position); + } + rowCount += positionCount - insertsFromUpdates; + } + + @Override + public void finish() + { + if (state == State.RUNNING) { + state = State.FINISHING; + finishFuture = toListenableFuture(mergeSink.finish()); + } + } + + @Override + protected void abort() + { + mergeSink.abort(); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java index 39f6a860ac1d5..29b66ba4b512a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java @@ -18,9 +18,12 @@ import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.OutputTableHandle; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.ConnectorMergeSink; import com.facebook.presto.spi.ConnectorPageSink; import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.PageSinkContext; +import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.connector.ConnectorPageSinkProvider; import java.util.concurrent.ConcurrentHashMap; @@ -73,6 +76,15 @@ public ConnectorPageSink createPageSink(Session session, InsertTableHandle table return createPageSink(session, tableHandle, pageSinkContext, null); } + @Override + public ConnectorMergeSink createMergeSink(Session session, MergeHandle mergeHandle) + { + // assumes connectorId and catalog are the same + TableHandle tableHandle = mergeHandle.getTableHandle(); + ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId()); + return providerFor(tableHandle.getConnectorId()).createMergeSink(tableHandle.getTransaction(), connectorSession, mergeHandle.getConnectorMergeTableHandle()); + } + private ConnectorPageSinkProvider providerFor(ConnectorId connectorId) { ConnectorPageSinkProvider provider = pageSinkProviders.get(connectorId); diff --git a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkProvider.java b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkProvider.java index 3e46127870194..25cb78cbe063b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkProvider.java +++ b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkProvider.java @@ -16,7 +16,9 @@ import com.facebook.presto.Session; import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.OutputTableHandle; +import com.facebook.presto.spi.ConnectorMergeSink; import com.facebook.presto.spi.ConnectorPageSink; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.PageSinkContext; public interface PageSinkProvider @@ -24,4 +26,9 @@ public interface PageSinkProvider ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext); ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext); + + /* + * Used to write the result of SQL MERGE to an existing table + */ + ConnectorMergeSink createMergeSink(Session session, MergeHandle mergeHandle); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java index eb2dc037dcf46..363c1222d13e3 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java @@ -33,6 +33,7 @@ import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.VarcharType; import com.facebook.presto.metadata.CatalogMetadata; +import com.facebook.presto.metadata.FunctionAndTypeManager; import com.facebook.presto.metadata.Metadata; import com.facebook.presto.metadata.OperatorNotFoundException; import com.facebook.presto.metadata.TableFunctionMetadata; @@ -41,6 +42,7 @@ import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.MaterializedViewDefinition; import com.facebook.presto.spi.MaterializedViewStatus; +import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.PrestoWarning; import com.facebook.presto.spi.SchemaTableName; @@ -70,6 +72,7 @@ import com.facebook.presto.spi.function.table.TableArgument; import com.facebook.presto.spi.function.table.TableArgumentSpecification; import com.facebook.presto.spi.function.table.TableFunctionAnalysis; +import com.facebook.presto.spi.plan.PartitioningHandle; import com.facebook.presto.spi.relation.DomainTranslator; import com.facebook.presto.spi.relation.RowExpression; import com.facebook.presto.spi.security.AccessControl; @@ -81,6 +84,7 @@ import com.facebook.presto.sql.ExpressionUtils; import com.facebook.presto.sql.MaterializedViewUtils; import com.facebook.presto.sql.SqlFormatterUtil; +import com.facebook.presto.sql.analyzer.Analysis.MergeAnalysis; import com.facebook.presto.sql.analyzer.Analysis.TableArgumentAnalysis; import com.facebook.presto.sql.analyzer.Analysis.TableFunctionInvocationAnalysis; import com.facebook.presto.sql.parser.ParsingException; @@ -146,6 +150,9 @@ import com.facebook.presto.sql.tree.LogicalBinaryExpression; import com.facebook.presto.sql.tree.LongLiteral; import com.facebook.presto.sql.tree.Merge; +import com.facebook.presto.sql.tree.MergeCase; +import com.facebook.presto.sql.tree.MergeInsert; +import com.facebook.presto.sql.tree.MergeUpdate; import com.facebook.presto.sql.tree.NaturalJoin; import com.facebook.presto.sql.tree.Node; import com.facebook.presto.sql.tree.NodeLocation; @@ -348,6 +355,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Throwables.throwIfInstanceOf; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; @@ -2061,6 +2069,16 @@ protected Scope visitTable(Table table, Optional scope) analysis.addSourceColumns(field, ImmutableSet.of(new SourceColumn(name, column.getName()))); } + boolean isMergeIntoStatement = statement instanceof Merge && ((Merge) statement).getTargetTable().equals(table); + if (isMergeIntoStatement) { + // Add the merge row id field + ColumnHandle mergeRowIdColumnHandle = metadata.getMergeRowIdColumnHandle(session, tableHandle.get()); + Type mergeRowIdType = metadata.getColumnMetadata(session, tableHandle.get(), mergeRowIdColumnHandle).getType(); + Field mergeRowIdField = Field.newUnqualified(Optional.empty(), Optional.empty(), mergeRowIdType); + fields.add(mergeRowIdField); + analysis.setColumn(mergeRowIdField, mergeRowIdColumnHandle); + } + analysis.registerTable(table, tableHandle.get()); List outputFields = fields.build(); @@ -2085,7 +2103,15 @@ protected Scope visitTable(Table table, Optional scope) } } - return createAndAssignScope(table, scope, outputFields); + Scope tableScope = createAndAssignScope(table, scope, outputFields); + + if (isMergeIntoStatement) { + FieldReference mergeRowIdFieldReference = new FieldReference(outputFields.size() - 1); + analyzeExpression(mergeRowIdFieldReference, tableScope); + analysis.setRowIdField(table, mergeRowIdFieldReference); + } + + return tableScope; } private Optional getTableHandle(TableColumnMetadata tableColumnsMetadata, Table table, QualifiedObjectName name, Optional scope) @@ -2974,7 +3000,260 @@ protected Scope visitUpdate(Update update, Optional scope) @Override protected Scope visitMerge(Merge merge, Optional scope) { - throw new PrestoException(StandardErrorCode.NOT_SUPPORTED, "This connector does not support MERGE INTO statements"); + Relation targetRelation = merge.getTarget(); + Table targetTable = getMergeTargetTable(targetRelation); + QualifiedObjectName targetTableQualifiedName = createQualifiedObjectName(session, targetTable, targetTable.getName(), metadata); + MetadataHandle metadataHandle = analysis.getMetadataHandle(); + + if (getViewDefinition(session, metadataResolver, metadataHandle, targetTableQualifiedName).isPresent()) { + throw new SemanticException(NOT_SUPPORTED, merge, "Merging into views is not supported"); + } + + if (getMaterializedViewDefinition(session, metadataResolver, metadataHandle, targetTableQualifiedName).isPresent()) { + throw new SemanticException(NOT_SUPPORTED, merge, "Merging into materialized views is not supported"); + } + + TableColumnMetadata targetTableColumnsMetadata = getTableColumnsMetadata(session, metadataResolver, metadataHandle, targetTableQualifiedName); + + TableHandle targetTableHandle = targetTableColumnsMetadata.getTableHandle() + .orElseThrow(() -> new SemanticException(MISSING_TABLE, targetTable, "Table '%s' does not exist", targetTableQualifiedName)); + + StatementAnalyzer statementAnalyzer = new StatementAnalyzer(analysis, metadata, sqlParser, + new AllowAllAccessControl(), session, warningCollector); + + Scope targetTableScope = statementAnalyzer.analyze(targetRelation, scope); + Scope sourceTableScope = process(merge.getSource(), scope); + Scope joinScope = createAndAssignScope(merge, scope, targetTableScope.getRelationType().joinWith(sourceTableScope.getRelationType())); + + List targetColumnsMetadata = targetTableColumnsMetadata.getColumnsMetadata().stream() + .filter(column -> !column.isHidden()) + .collect(toImmutableList()); + + Optional targetInsertLayout = metadata.getInsertLayout(session, targetTableHandle); + + Map targetAllColumnHandles = metadata.getColumnHandles(session, targetTableHandle); + ImmutableList.Builder targetColumnHandlesBuilder = ImmutableList.builder(); + ImmutableSet.Builder targetColumnNamesBuilder = ImmutableSet.builder(); + ImmutableList.Builder targetRedistributionColumnHandlesBuilder = ImmutableList.builder(); + Set targetPartitioningColumnNames = ImmutableSet.copyOf(targetInsertLayout.map(NewTableLayout::getPartitionColumns).orElse(ImmutableList.of())); + for (ColumnMetadata columnMetadata : targetColumnsMetadata) { + String targetColumnName = columnMetadata.getName(); + ColumnHandle targetColumnHandle = targetAllColumnHandles.get(targetColumnName); + targetColumnHandlesBuilder.add(targetColumnHandle); + targetColumnNamesBuilder.add(targetColumnName); + if (targetPartitioningColumnNames.contains(targetColumnName)) { + targetRedistributionColumnHandlesBuilder.add(targetColumnHandle); + } + } + List targetColumnHandles = targetColumnHandlesBuilder.build(); + Set targetColumnNames = targetColumnNamesBuilder.build(); + // The "targetRedistributionColumnHandles" is a list that contains the columns from the target table that are also present in the partitioning columns. + List targetRedistributionColumnHandles = targetRedistributionColumnHandlesBuilder.build(); + + Map targetColumnTypes = targetColumnsMetadata.stream().collect(toImmutableMap(ColumnMetadata::getName, ColumnMetadata::getType)); + + // Analyze all expressions in the Merge node + + Expression mergePredicate = merge.getPredicate(); + ExpressionAnalysis mergePredicateAnalysis = analyzeExpression(mergePredicate, joinScope); + Type mergePredicateType = mergePredicateAnalysis.getType(mergePredicate); + if (!mergePredicateType.equals(BOOLEAN)) { + if (!mergePredicateType.equals(UNKNOWN)) { + throw new SemanticException(TYPE_MISMATCH, mergePredicate, "The MERGE predicate must evaluate to a boolean: actual type %s", mergePredicateType); + } + // coerce null to boolean + analysis.addCoercion(mergePredicate, BOOLEAN, false); + } + analysis.recordSubqueries(merge, mergePredicateAnalysis); + + Set allUpdateColumnNames = new HashSet<>(); + + for (int caseCounter = 0; caseCounter < merge.getMergeCases().size(); caseCounter++) { + MergeCase mergeCase = merge.getMergeCases().get(caseCounter); + List setColumnNames = lowercaseIdentifierList(mergeCase.getSetColumns()); + if (mergeCase instanceof MergeUpdate) { + allUpdateColumnNames.addAll(setColumnNames); + } + else if (mergeCase instanceof MergeInsert && setColumnNames.isEmpty()) { + setColumnNames = targetColumnsMetadata.stream().map(ColumnMetadata::getName).collect(toImmutableList()); + } + int mergeCaseSetColumnCount = setColumnNames.size(); + List mergeCaseSetExpressions = mergeCase.getSetExpressions(); + checkArgument( + mergeCaseSetColumnCount == mergeCaseSetExpressions.size(), + "Number of merge columns (%s) isn't equal to number of expressions (%s)", + mergeCaseSetColumnCount, mergeCaseSetExpressions.size()); + Set mergeCaseColumnNameSet = new HashSet<>(mergeCaseSetColumnCount); + // Look for missing or duplicate column names. + setColumnNames.forEach(mergeCaseColumnName -> { + if (!targetColumnNames.contains(mergeCaseColumnName)) { + throw new SemanticException(MISSING_COLUMN, merge, "Merge column name does not exist in target table: %s", mergeCaseColumnName); + } + if (!mergeCaseColumnNameSet.add(mergeCaseColumnName)) { + throw new SemanticException(DUPLICATE_COLUMN_NAME, merge, "Merge column name is specified more than once: %s", mergeCaseColumnName); + } + }); + + // Collects types for columns and expressions in this MergeCase. + ImmutableList.Builder setColumnTypesBuilder = ImmutableList.builder(); + ImmutableList.Builder setExpressionTypesBuilder = ImmutableList.builder(); + for (int index = 0; index < setColumnNames.size(); index++) { + String columnName = setColumnNames.get(index); + Expression setExpression = mergeCaseSetExpressions.get(index); + ExpressionAnalysis setExpressionAnalysis = analyzeExpression(setExpression, joinScope); + analysis.recordSubqueries(merge, setExpressionAnalysis); + Type setColumnType = requireNonNull(targetColumnTypes.get(columnName)); + setColumnTypesBuilder.add(setColumnType); + setExpressionTypesBuilder.add(setExpressionAnalysis.getType(setExpression)); + } + List setColumnTypes = setColumnTypesBuilder.build(); + List setExpressionTypes = setExpressionTypesBuilder.build(); + + // Check if the types of the columns and expressions match for the MERGE SET clause. + if (!checkTypesMatchForMergeSet(setColumnTypes, setExpressionTypes)) { + throw new SemanticException(TYPE_MISMATCH, + mergeCase, + "MERGE table column types don't match for MERGE case %s, SET expressions: Table: [%s], Expressions: [%s]", + caseCounter, + Joiner.on(", ").join(setColumnTypes), + Joiner.on(", ").join(setExpressionTypes)); + } + + // Add coercion if the target column type and set expression type do not match. + for (int index = 0; index < setColumnNames.size(); index++) { + Expression setExpression = mergeCase.getSetExpressions().get(index); + Type targetColumnType = targetColumnTypes.get(setColumnNames.get(index)); + Type setExpressionType = setExpressionTypes.get(index); + if (!targetColumnType.equals(setExpressionType)) { + FunctionAndTypeManager functionAndTypeManager = metadata.getFunctionAndTypeManager(); + analysis.addCoercion(setExpression, targetColumnType, functionAndTypeManager.isTypeOnlyCoercion(setExpressionType, targetColumnType)); + } + } + } + + // Check if the user has permission to insert into the target table + merge.getMergeCases().stream() + .filter(mergeCase -> mergeCase instanceof MergeInsert) + .findFirst() + .ifPresent(mergeCase -> accessControl.checkCanInsertIntoTable(session.getRequiredTransactionId(), + session.getIdentity(), session.getAccessControlContext(), targetTableQualifiedName)); + + // If there are any columns to update then verify the user has permission to update these columns. + if (!allUpdateColumnNames.isEmpty()) { + accessControl.checkCanUpdateTableColumns(session.getRequiredTransactionId(), session.getIdentity(), + session.getAccessControlContext(), targetTableQualifiedName, allUpdateColumnNames); + } + + analysis.setUpdateType("MERGE"); + + List> mergeCaseColumnHandles = buildMergeCaseColumnLists(merge, targetColumnsMetadata, targetAllColumnHandles); + + Optional mergeUpdateLayout = metadata.getMergeUpdateLayout(session, targetTableHandle); + + ImmutableMap.Builder columnHandleFieldNumbersBuilder = ImmutableMap.builder(); + Map fieldIndexes = new HashMap<>(); + RelationType targetRelationType = targetTableScope.getRelationType(); + for (Field targetField : targetRelationType.getAllFields()) { + // Only the rowId column handle will have no name, and we want to skip that column + targetField.getName().ifPresent(targetFieldName -> { + int targetFieldIndex = targetRelationType.indexOf(targetField); + ColumnHandle targetColumnHandle = targetAllColumnHandles.get(targetFieldName); + verify(targetColumnHandle != null, "targetAllColumnHandles does not contain the named handle: %s", targetFieldName); + columnHandleFieldNumbersBuilder.put(targetColumnHandle, targetFieldIndex); + fieldIndexes.put(targetFieldName, targetFieldIndex); + }); + } + Map columnHandleFieldNumbers = columnHandleFieldNumbersBuilder.buildOrThrow(); + + List insertPartitioningArgumentIndexes = targetPartitioningColumnNames.stream() + .map(fieldIndexes::get) + .collect(toImmutableList()); + + Set nonNullableColumnHandles = metadata.getTableMetadata(session, targetTableHandle).getColumns().stream() + .filter(column -> !column.isNullable()) + .map(ColumnMetadata::getName) + .map(targetAllColumnHandles::get) + .collect(toImmutableSet()); + + analysis.setMergeAnalysis(new MergeAnalysis( + targetTable, + targetColumnsMetadata, + targetColumnHandles, + targetRedistributionColumnHandles, + mergeCaseColumnHandles, + nonNullableColumnHandles, + columnHandleFieldNumbers, + insertPartitioningArgumentIndexes, + targetInsertLayout, + mergeUpdateLayout, + targetTableScope, + joinScope)); + + return createAndAssignScope(merge, Optional.empty(), Field.newUnqualified(merge.getLocation(), "rows", BIGINT)); + } + + private boolean checkTypesMatchForMergeSet(Iterable tableTypes, Iterable queryTypes) + { + if (Iterables.size(tableTypes) != Iterables.size(queryTypes)) { + return false; + } + + Iterator tableTypesIterator = tableTypes.iterator(); + Iterator queryTypesIterator = queryTypes.iterator(); + while (tableTypesIterator.hasNext()) { + Type tableType = tableTypesIterator.next(); + Type queryType = queryTypesIterator.next(); + + if (!metadata.getFunctionAndTypeManager().canCoerce(queryType, tableType)) { + return false; + } + } + + return true; + } + + private Table getMergeTargetTable(Relation relation) + { + if (relation instanceof Table) { + return (Table) relation; + } + checkArgument(relation instanceof AliasedRelation, "relation is neither a Table nor an AliasedRelation"); + return (Table) ((AliasedRelation) relation).getRelation(); + } + + /** + * Builds a list of column handles for each merge case in the given merge statement. + * + * @param merge the merge statement + * @param columnSchemas the list of column metadata for the target table. + * @param allColumnHandles a map of column names to column handles for the target table. + * @return a list of lists of column handles, where each inner list corresponds to a merge case. + */ + private List> buildMergeCaseColumnLists(Merge merge, List columnSchemas, Map allColumnHandles) + { + ImmutableList.Builder> mergeCaseColumnsListsBuilder = ImmutableList.builder(); + for (int caseCounter = 0; caseCounter < merge.getMergeCases().size(); caseCounter++) { + MergeCase mergeCase = merge.getMergeCases().get(caseCounter); + List mergeColumnNames; + if (mergeCase instanceof MergeInsert && mergeCase.getSetColumns().isEmpty()) { + mergeColumnNames = columnSchemas.stream().map(ColumnMetadata::getName).collect(toImmutableList()); + } + else { + mergeColumnNames = lowercaseIdentifierList(mergeCase.getSetColumns()); + } + mergeCaseColumnsListsBuilder.add( + mergeColumnNames.stream() + .map(name -> requireNonNull(allColumnHandles.get(name), "No column found for name")) + .collect(toImmutableList())); + } + return mergeCaseColumnsListsBuilder.build(); + } + + private List lowercaseIdentifierList(Collection identifiers) + { + return identifiers.stream() + .map(identifier -> identifier.getValue().toLowerCase(ENGLISH)) + .collect(toImmutableList()); } private Scope analyzeJoinUsing(Join node, List columns, Optional scope, Scope left, Scope right) diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/BasePlanFragmenter.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/BasePlanFragmenter.java index bc77373a5ecbd..62969282c375c 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/BasePlanFragmenter.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/BasePlanFragmenter.java @@ -43,6 +43,8 @@ import com.facebook.presto.spi.relation.VariableReferenceExpression; import com.facebook.presto.sql.planner.plan.ExchangeNode; import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.SimplePlanRewriter; @@ -263,6 +265,21 @@ public PlanNode visitTableWriter(TableWriterNode node, RewriteContext context) + { + if (node.getPartitioningScheme().isPresent()) { + context.get().setDistribution(node.getPartitioningScheme().get().getPartitioning().getHandle(), metadata, session); + } + return context.defaultRewrite(node, context.get()); + } + + @Override + public PlanNode visitMergeProcessor(MergeProcessorNode node, RewriteContext context) + { + return context.defaultRewrite(node, context.get()); + } + @Override public PlanNode visitValues(ValuesNode node, RewriteContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java index b5a460f033370..676687c0f1c6a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java @@ -37,6 +37,7 @@ import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.CreateHandle; import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.DeleteHandle; import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.InsertHandle; +import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.MergeHandle; import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.RefreshMaterializedViewHandle; import com.facebook.presto.execution.scheduler.ExecutionWriterTarget.UpdateHandle; import com.facebook.presto.execution.scheduler.TableWriteInfo; @@ -73,6 +74,8 @@ import com.facebook.presto.operator.LookupOuterOperator.LookupOuterOperatorFactory; import com.facebook.presto.operator.LookupSourceFactory; import com.facebook.presto.operator.MarkDistinctOperator.MarkDistinctOperatorFactory; +import com.facebook.presto.operator.MergeProcessorOperator; +import com.facebook.presto.operator.MergeWriterOperator; import com.facebook.presto.operator.MetadataDeleteOperator.MetadataDeleteOperatorFactory; import com.facebook.presto.operator.NestedLoopJoinBridge; import com.facebook.presto.operator.NestedLoopJoinPagesSupplier; @@ -206,6 +209,8 @@ import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; @@ -2768,6 +2773,47 @@ private PageSinkCommitStrategy getPageSinkCommitStrategy() return NO_COMMIT; } + @Override + public PhysicalOperation visitMergeWriter(MergeWriterNode node, LocalExecutionPlanContext context) + { + context.setDriverInstanceCount(getTaskWriterCount(session)); + + PhysicalOperation source = node.getSource().accept(this, context); + OperatorFactory operatorFactory = new MergeWriterOperator.MergeWriterOperatorFactory( + context.getNextOperatorId(), node.getId(), pageSinkManager, node.getTarget(), session, + tableCommitContextCodec); + return new PhysicalOperation(operatorFactory, makeLayout(node), context, source); + } + + @Override + public PhysicalOperation visitMergeProcessor(MergeProcessorNode node, LocalExecutionPlanContext context) + { + PhysicalOperation source = node.getSource().accept(this, context); + + Map nodeLayout = makeLayout(node); + Map sourceLayout = makeLayout(node.getSource()); + int rowIdChannel = sourceLayout.get(node.getRowIdVariable()); + int mergeRowChannel = sourceLayout.get(node.getMergeRowVariable()); + + List redistributionColumns = node.getTargetRedistributionColumnVariables().stream() + .map(nodeLayout::get) + .collect(toImmutableList()); + List targetColumnChannels = node.getTargetColumnVariables().stream() + .map(nodeLayout::get) + .collect(toImmutableList()); + + OperatorFactory operatorFactory = new MergeProcessorOperator.MergeProcessorOperatorFactory( + context.getNextOperatorId(), + node.getId(), + node.getTarget().getMergeParadigmAndTypes(), + rowIdChannel, + mergeRowChannel, + redistributionColumns, + targetColumnChannels); + + return new PhysicalOperation(operatorFactory, nodeLayout, context, source); + } + @Override public PhysicalOperation visitStatisticsWriterNode(StatisticsWriterNode node, LocalExecutionPlanContext context) { @@ -3494,6 +3540,10 @@ else if (target instanceof UpdateHandle) { metadata.finishUpdate(session, ((UpdateHandle) target).getHandle(), fragments); return Optional.empty(); } + else if (target instanceof MergeHandle) { + metadata.finishMerge(session, ((MergeHandle) target).getHandle(), fragments, statistics); + return Optional.empty(); + } else { throw new AssertionError("Unhandled target type: " + target.getClass().getName()); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java index e1c1979f924be..4f4fa957220b4 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java @@ -54,6 +54,7 @@ import com.facebook.presto.sql.parser.SqlParser; import com.facebook.presto.sql.planner.StatisticsAggregationPlanner.TableStatisticAggregation; import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; import com.facebook.presto.sql.planner.plan.UpdateNode; import com.facebook.presto.sql.tree.Analyze; @@ -66,6 +67,7 @@ import com.facebook.presto.sql.tree.Identifier; import com.facebook.presto.sql.tree.Insert; import com.facebook.presto.sql.tree.LambdaArgumentDeclaration; +import com.facebook.presto.sql.tree.Merge; import com.facebook.presto.sql.tree.NodeRef; import com.facebook.presto.sql.tree.NullLiteral; import com.facebook.presto.sql.tree.Parameter; @@ -103,6 +105,7 @@ import static com.facebook.presto.sql.analyzer.ExpressionTreeUtils.createSymbolReference; import static com.facebook.presto.sql.analyzer.ExpressionTreeUtils.getSourceLocation; import static com.facebook.presto.sql.planner.PlannerUtils.newVariable; +import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION; import static com.facebook.presto.sql.planner.TranslateExpressionsUtil.toRowExpression; import static com.facebook.presto.sql.relational.Expressions.constant; import static com.facebook.presto.sql.tree.ExplainFormat.Type.TEXT; @@ -181,6 +184,9 @@ else if (statement instanceof Delete) { if (statement instanceof Update) { return createUpdatePlan(analysis, (Update) statement); } + if (statement instanceof Merge) { + return createMergePlan(analysis, (Merge) statement); + } else if (statement instanceof Query) { return createRelationPlan(analysis, (Query) statement, new SqlPlannerContext(0)); } @@ -537,6 +543,26 @@ private RelationPlan createUpdatePlan(Analysis analysis, Update node) return new RelationPlan(commitNode, analysis.getScope(node), commitNode.getOutputVariables()); } + private RelationPlan createMergePlan(Analysis analysis, Merge node) + { + SqlPlannerContext context = new SqlPlannerContext(0); + MergeWriterNode mergeNode = new QueryPlanner(analysis, variableAllocator, idAllocator, + buildLambdaDeclarationToVariableMap(analysis, variableAllocator), metadata, session, context, sqlParser) + .plan(node); + + TableFinishNode commitNode = new TableFinishNode( + mergeNode.getSourceLocation(), + idAllocator.getNextId(), + mergeNode, + Optional.of(mergeNode.getTarget()), + variableAllocator.newVariable("rows", BIGINT), + Optional.empty(), + Optional.empty(), + Optional.empty()); + + return new RelationPlan(commitNode, analysis.getScope(node), commitNode.getOutputVariables()); + } + private PlanNode createOutputPlan(RelationPlan plan, Analysis analysis) { ImmutableList.Builder outputs = ImmutableList.builder(); @@ -640,7 +666,7 @@ private static Optional getPartitioningSchemeForTableWrite(O List outputLayout = new ArrayList<>(variables); partitioningScheme = Optional.of(new PartitioningScheme( - Partitioning.create(tableLayout.get().getPartitioning(), partitionFunctionArguments), + Partitioning.create(tableLayout.get().getPartitioning().orElse(FIXED_HASH_DISTRIBUTION), partitionFunctionArguments), outputLayout, tableLayout.get().getWriterPolicy() == MULTIPLE_WRITERS_PER_PARTITION_ALLOWED)); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/MergePartitioningHandle.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/MergePartitioningHandle.java new file mode 100644 index 0000000000000..0dfd0c83793ee --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/MergePartitioningHandle.java @@ -0,0 +1,210 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner; + +import com.facebook.presto.common.Page; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.common.type.Type; +import com.facebook.presto.operator.BucketPartitionFunction; +import com.facebook.presto.operator.PartitionFunction; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; +import com.facebook.presto.spi.plan.PartitioningHandle; +import com.facebook.presto.spi.plan.PartitioningScheme; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.VerifyException; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.IntStream; + +import static com.facebook.presto.common.type.TinyintType.TINYINT; +import static com.facebook.presto.spi.ConnectorMergeSink.DELETE_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.INSERT_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.UPDATE_OPERATION_NUMBER; +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.Iterables.getLast; +import static java.util.Objects.requireNonNull; + +public final class MergePartitioningHandle + implements ConnectorPartitioningHandle +{ + private final Optional insertPartitioning; + private final Optional updatePartitioning; + + @JsonCreator + public MergePartitioningHandle( + @JsonProperty("insertPartitioning") Optional insertPartitioning, + @JsonProperty("updatePartitioning") Optional updatePartitioning) + { + this.insertPartitioning = requireNonNull(insertPartitioning, "insertPartitioning is null"); + this.updatePartitioning = requireNonNull(updatePartitioning, "updatePartitioning is null"); + checkArgument(insertPartitioning.isPresent() || updatePartitioning.isPresent(), "insert or update partitioning must be present"); + } + + @JsonProperty("insertPartitioning") + public Optional getInsertPartitioning() + { + return insertPartitioning; + } + + @JsonProperty("updatePartitioning") + public Optional getUpdatePartitioning() + { + return updatePartitioning; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + MergePartitioningHandle that = (MergePartitioningHandle) o; + return insertPartitioning.equals(that.insertPartitioning) && + updatePartitioning.equals(that.updatePartitioning); + } + + @Override + public int hashCode() + { + return Objects.hash(insertPartitioning, updatePartitioning); + } + + @Override + public String toString() + { + List parts = new ArrayList<>(); + insertPartitioning.ifPresent(scheme -> parts.add("insert = " + scheme.getPartitioning().getHandle())); + updatePartitioning.ifPresent(scheme -> parts.add("update = " + scheme.getPartitioning().getHandle())); + return "MERGE " + parts; + } + + public NodePartitionMap getNodePartitioningMap(Function getMap) + { + Optional optionalInsertMap = insertPartitioning.map(scheme -> scheme.getPartitioning().getHandle()).map(getMap); + Optional optionalUpdateMap = updatePartitioning.map(scheme -> scheme.getPartitioning().getHandle()).map(getMap); + + if (optionalInsertMap.isPresent() && optionalUpdateMap.isPresent()) { + NodePartitionMap insertMap = optionalInsertMap.get(); + NodePartitionMap updateMap = optionalUpdateMap.get(); + if (!insertMap.getPartitionToNode().equals(updateMap.getPartitionToNode()) || + !Arrays.equals(insertMap.getBucketToPartition(), updateMap.getBucketToPartition())) { + throw new PrestoException(NOT_SUPPORTED, "Insert and update layout have mismatched BucketNodeMap"); + } + } + + return optionalInsertMap.orElseGet(optionalUpdateMap::get); + } + + public PartitionFunction getPartitionFunction(PartitionFunctionLookup partitionFunctionLookup, List types, int[] bucketToPartition) + { + // channels: merge row, insert arguments, update row ID + List insertTypes = types.subList(1, types.size() - (updatePartitioning.isPresent() ? 1 : 0)); + + Optional insertFunction = insertPartitioning.map(scheme -> + partitionFunctionLookup.get(scheme, insertTypes)); + + Optional updateFunction = updatePartitioning.map(scheme -> + partitionFunctionLookup.get(scheme, Collections.singletonList(getLast(types)))); + + return getPartitionFunction(insertFunction, updateFunction, insertTypes.size(), bucketToPartition); + } + + private static PartitionFunction getPartitionFunction(Optional insertFunction, Optional updateFunction, int insertArguments, int[] bucketToPartition) + { + if (insertFunction.isPresent() && updateFunction.isPresent()) { + return new MergePartitionFunction( + insertFunction.get(), + updateFunction.get(), + IntStream.range(1, insertArguments + 1).toArray(), + new int[] {insertArguments + 1}); + } + + PartitionFunction roundRobinFunction = new BucketPartitionFunction(new SystemPartitioningHandle.SystemPartitionFunction.RoundRobinBucketFunction(bucketToPartition.length), bucketToPartition); + + if (insertFunction.isPresent()) { + return new MergePartitionFunction( + insertFunction.get(), + roundRobinFunction, + IntStream.range(1, insertArguments + 1).toArray(), + new int[] {}); + } + + if (updateFunction.isPresent()) { + return new MergePartitionFunction( + roundRobinFunction, + updateFunction.get(), + new int[] {}, + new int[] {insertArguments + 1}); + } + + throw new AssertionError(); + } + + public interface PartitionFunctionLookup + { + PartitionFunction get(PartitioningScheme scheme, List partitionChannelTypes); + } + + private static final class MergePartitionFunction + implements PartitionFunction + { + private final PartitionFunction insertFunction; + private final PartitionFunction updateFunction; + private final int[] insertColumns; + private final int[] updateColumns; + + public MergePartitionFunction(PartitionFunction insertFunction, PartitionFunction updateFunction, int[] insertColumns, int[] updateColumns) + { + this.insertFunction = requireNonNull(insertFunction, "insertFunction is null"); + this.updateFunction = requireNonNull(updateFunction, "updateFunction is null"); + this.insertColumns = requireNonNull(insertColumns, "insertColumns is null"); + this.updateColumns = requireNonNull(updateColumns, "updateColumns is null"); + checkArgument(insertFunction.getPartitionCount() == updateFunction.getPartitionCount(), "partition counts must match"); + } + + @Override + public int getPartitionCount() + { + return insertFunction.getPartitionCount(); + } + + @Override + public int getPartition(Page page, int position) + { + Block operationBlock = page.getBlock(0); + byte operation = TINYINT.getByte(operationBlock, position); + switch (operation) { + case INSERT_OPERATION_NUMBER: + return insertFunction.getPartition(page.getColumns(insertColumns), position); + case UPDATE_OPERATION_NUMBER: + case DELETE_OPERATION_NUMBER: + return updateFunction.getPartition(page.getColumns(updateColumns), position); + default: + throw new VerifyException("Invalid merge operation number: " + operation); + } + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/NodePartitioningManager.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/NodePartitioningManager.java index 3de6783a91a41..7806e1f7a9843 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/NodePartitioningManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/NodePartitioningManager.java @@ -20,6 +20,7 @@ import com.facebook.presto.execution.scheduler.NodeScheduler; import com.facebook.presto.execution.scheduler.group.DynamicBucketNodeMap; import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelectionStats; +import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelector; import com.facebook.presto.metadata.InternalNode; import com.facebook.presto.metadata.Split; import com.facebook.presto.operator.BucketPartitionFunction; @@ -36,6 +37,7 @@ import com.facebook.presto.spi.plan.PartitioningScheme; import com.facebook.presto.spi.schedule.NodeSelectionStrategy; import com.facebook.presto.split.EmptySplit; +import com.facebook.presto.sql.planner.SystemPartitioningHandle.SystemPartitioning; import com.google.common.collect.BiMap; import com.google.common.collect.HashBiMap; import com.google.common.collect.ImmutableList; @@ -49,11 +51,16 @@ import java.util.function.ToIntFunction; import java.util.stream.IntStream; +import static com.facebook.presto.SystemSessionProperties.getHashPartitionCount; import static com.facebook.presto.SystemSessionProperties.getMaxTasksPerStage; import static com.facebook.presto.metadata.InternalNode.NodeStatus.DEAD; import static com.facebook.presto.spi.StandardErrorCode.NODE_SELECTION_NOT_SUPPORTED; +import static com.facebook.presto.spi.StandardErrorCode.NO_NODES_AVAILABLE; +import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION; +import static com.facebook.presto.util.Failures.checkCondition; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.Math.min; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -89,6 +96,12 @@ public PartitionFunction getPartitionFunction( partitioningScheme.getHashColumn().isPresent(), partitioningScheme.getBucketToPartition().get()); } + else if (partitioningHandle.getConnectorHandle() instanceof MergePartitioningHandle) { + MergePartitioningHandle handle = (MergePartitioningHandle) partitioningHandle.getConnectorHandle(); + return handle.getPartitionFunction( + (scheme, types) -> getPartitionFunction(session, scheme, types), + partitionChannelTypes, bucketToPartition.get()); + } else { ConnectorNodePartitioningProvider partitioningProvider = partitioningProviderManager.getPartitioningProvider(partitioningHandle.getConnectorId().get()); @@ -120,18 +133,35 @@ public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHand return getNodePartitioningMap(session, partitioningHandle, Optional.empty()); } + /** + * This method is recursive for MergePartitioningHandle. It caches the node mappings + * to ensure that both the insert and update layouts use the same mapping. + */ public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHandle partitioningHandle, Optional> nodePredicate) { requireNonNull(session, "session is null"); requireNonNull(partitioningHandle, "partitioningHandle is null"); if (partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle) { - return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getNodePartitionMap(session, nodeScheduler, nodePredicate); + // TODO #20578: The next commented line is the original code. The following one is an alternative that I need to check if is valid. + // return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getNodePartitionMap(session, nodeScheduler, nodePredicate); + return systemNodePartitionMap(session, partitioningHandle, nodePredicate); + } + + if (partitioningHandle.getConnectorHandle() instanceof MergePartitioningHandle) { + MergePartitioningHandle mergeHandle = (MergePartitioningHandle) partitioningHandle.getConnectorHandle(); + return mergeHandle.getNodePartitioningMap(handle -> getNodePartitioningMap(session, handle, nodePredicate)); } ConnectorId connectorId = partitioningHandle.getConnectorId() .orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle)); - ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle, nodePredicate); + + Optional optionalMap = getConnectorBucketNodeMap(session, partitioningHandle, nodePredicate); + if (!optionalMap.isPresent()) { + return systemNodePartitionMap(session, FIXED_HASH_DISTRIBUTION, nodePredicate); + } + ConnectorBucketNodeMap connectorBucketNodeMap = optionalMap.get(); + // safety check for crazy partitioning checkArgument(connectorBucketNodeMap.getBucketCount() < 1_000_000, "Too many buckets in partitioning: %s", connectorBucketNodeMap.getBucketCount()); @@ -179,32 +209,58 @@ public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHand public BucketNodeMap getBucketNodeMap(Session session, PartitioningHandle partitioningHandle, boolean preferDynamic) { - ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle, Optional.empty()); + Optional connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle, Optional.empty()); - NodeSelectionStrategy nodeSelectionStrategy = connectorBucketNodeMap.getNodeSelectionStrategy(); + int bucketCount = getBucketCount(session, partitioningHandle, connectorBucketNodeMap, preferDynamic); + + // TODO #20578: WIP - This method is under development. Unsafe ".get()" method call. + NodeSelectionStrategy nodeSelectionStrategy = connectorBucketNodeMap.get().getNodeSelectionStrategy(); switch (nodeSelectionStrategy) { case HARD_AFFINITY: - return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap), false); + return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap.get()), false); case SOFT_AFFINITY: if (preferDynamic) { - return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), connectorBucketNodeMap.getBucketCount(), getFixedMapping(connectorBucketNodeMap)); + return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), bucketCount, getFixedMapping(connectorBucketNodeMap.get())); } - return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap), true); + return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap.get()), true); case NO_PREFERENCE: if (preferDynamic) { - return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), connectorBucketNodeMap.getBucketCount()); + return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), bucketCount); } + return new FixedBucketNodeMap( getSplitToBucket(session, partitioningHandle), createArbitraryBucketToNode( nodeScheduler.createNodeSelector(session, partitioningHandle.getConnectorId().get()).selectRandomNodes(getMaxTasksPerStage(session)), - connectorBucketNodeMap.getBucketCount()), + bucketCount), false); default: throw new PrestoException(NODE_SELECTION_NOT_SUPPORTED, format("Unsupported node selection strategy %s", nodeSelectionStrategy)); } } + private Integer getBucketCount( + Session session, + PartitioningHandle partitioningHandle, + Optional connectorBucketNodeMap, + boolean preferDynamic) + { + Optional bucketCount = connectorBucketNodeMap.map(ConnectorBucketNodeMap::getBucketCount); + + return bucketCount.orElseGet(() -> preferDynamic ? + getNodeCount(session, partitioningHandle) : getAllNodes(session, partitioningHandle).size()); + } + + public int getNodeCount(Session session, PartitioningHandle partitioningHandle) + { + return getAllNodes(session, partitioningHandle).size(); + } + + private List getAllNodes(Session session, PartitioningHandle partitioningHandle) + { + return nodeScheduler.createNodeSelector(session, partitioningHandle.getConnectorId().get()).getAllNodes(); + } + private static List getFixedMapping(ConnectorBucketNodeMap connectorBucketNodeMap) { return connectorBucketNodeMap.getFixedMapping().stream() @@ -212,7 +268,7 @@ private static List getFixedMapping(ConnectorBucketNodeMap connect .collect(toImmutableList()); } - private ConnectorBucketNodeMap getConnectorBucketNodeMap(Session session, PartitioningHandle partitioningHandle, Optional> nodePredicate) + private Optional getConnectorBucketNodeMap(Session session, PartitioningHandle partitioningHandle, Optional> nodePredicate) { checkArgument(!(partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle)); ConnectorId connectorId = partitioningHandle.getConnectorId() @@ -221,14 +277,11 @@ private ConnectorBucketNodeMap getConnectorBucketNodeMap(Session session, Partit ConnectorNodePartitioningProvider partitioningProvider = partitioningProviderManager.getPartitioningProvider(partitioningHandle.getConnectorId().get()); - ConnectorBucketNodeMap connectorBucketNodeMap = partitioningProvider.getBucketNodeMap( + return partitioningProvider.getBucketNodeMap( partitioningHandle.getTransactionHandle().orElse(null), session.toConnectorSession(partitioningHandle.getConnectorId().get()), partitioningHandle.getConnectorHandle(), nodes); - - checkArgument(connectorBucketNodeMap != null, "No partition map %s", partitioningHandle); - return connectorBucketNodeMap; } private ToIntFunction getSplitToBucket(Session session, PartitioningHandle partitioningHandle) @@ -290,4 +343,31 @@ public List getNodes(Session session, ConnectorId connectorId, Optional> nodePredicate) + { + SystemPartitioning partitioning = ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getPartitioning(); + + // TODO #20578: Be careful with the "partitioningHandle.getConnectorId().orElse(null)". Evaluate possible side effects. + NodeSelector nodeSelector = nodeScheduler.createNodeSelector(session, partitioningHandle.getConnectorId().orElse(null), nodePredicate); + + List nodes; + if (partitioning == SystemPartitioning.COORDINATOR_ONLY) { + nodes = ImmutableList.of(nodeSelector.selectCurrentNode()); + } + else if (partitioning == SystemPartitioning.SINGLE) { + nodes = nodeSelector.selectRandomNodes(1); + } + else if (partitioning == SystemPartitioning.FIXED) { + nodes = nodeSelector.selectRandomNodes(min(getHashPartitionCount(session), getMaxTasksPerStage(session))); + } + else { + throw new IllegalArgumentException("Unsupported plan distribution " + partitioning); + } + checkCondition(!nodes.isEmpty(), NO_NODES_AVAILABLE, "No worker nodes available"); + + return new NodePartitionMap(nodes, split -> { + throw new UnsupportedOperationException("System distribution does not support source splits"); + }); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanBuilder.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanBuilder.java index 8729e4831c9e2..25cd1d1d8b57d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanBuilder.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanBuilder.java @@ -25,6 +25,8 @@ import com.facebook.presto.sql.analyzer.Analysis; import com.facebook.presto.sql.parser.SqlParser; import com.facebook.presto.sql.tree.Expression; +import com.facebook.presto.sql.tree.LambdaArgumentDeclaration; +import com.facebook.presto.sql.tree.NodeRef; import com.google.common.collect.ImmutableMap; import java.util.Map; @@ -47,6 +49,16 @@ public PlanBuilder(TranslationMap translations, PlanNode root) this.root = root; } + public static PlanBuilder newPlanBuilder( + RelationPlan plan, + Analysis analysis, + Map, VariableReferenceExpression> lambdaArguments) + { + return new PlanBuilder( + new TranslationMap(plan, analysis, lambdaArguments), + plan.getRoot()); + } + public TranslationMap copyTranslations() { TranslationMap translations = new TranslationMap(getRelationPlan(), getAnalysis(), getTranslations().getLambdaDeclarationToVariableMap()); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java index 2d7c8be053645..516b3a29f10c7 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java @@ -75,6 +75,7 @@ import com.facebook.presto.sql.planner.iterative.rule.PruneJoinColumns; import com.facebook.presto.sql.planner.iterative.rule.PruneLimitColumns; import com.facebook.presto.sql.planner.iterative.rule.PruneMarkDistinctColumns; +import com.facebook.presto.sql.planner.iterative.rule.PruneMergeSourceColumns; import com.facebook.presto.sql.planner.iterative.rule.PruneOrderByInAggregation; import com.facebook.presto.sql.planner.iterative.rule.PruneOutputColumns; import com.facebook.presto.sql.planner.iterative.rule.PruneProjectColumns; @@ -302,6 +303,7 @@ public PlanOptimizers( new PruneJoinColumns(), new PruneUpdateSourceColumns(), new PruneMarkDistinctColumns(), + new PruneMergeSourceColumns(), // TODO #20578: Check if this optimization still necessary. Is pruning columns with PruneUnreferencedOutputs enough? new PruneOutputColumns(), new PruneProjectColumns(), new PruneSemiJoinColumns(), diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java index 10ad2b1f68e7d..d388d2c43f911 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java @@ -18,13 +18,18 @@ import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.block.SortOrder; import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.common.type.RowType; import com.facebook.presto.common.type.Type; import com.facebook.presto.metadata.Metadata; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; +import com.facebook.presto.spi.MergeHandle; +import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.TableHandle; +import com.facebook.presto.spi.TableMetadata; import com.facebook.presto.spi.VariableAllocator; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.function.FunctionHandle; import com.facebook.presto.spi.function.FunctionMetadata; import com.facebook.presto.spi.plan.AggregationNode; @@ -34,14 +39,19 @@ import com.facebook.presto.spi.plan.DeleteNode; import com.facebook.presto.spi.plan.FilterNode; import com.facebook.presto.spi.plan.LimitNode; +import com.facebook.presto.spi.plan.MarkDistinctNode; import com.facebook.presto.spi.plan.Ordering; import com.facebook.presto.spi.plan.OrderingScheme; +import com.facebook.presto.spi.plan.Partitioning; +import com.facebook.presto.spi.plan.PartitioningHandle; +import com.facebook.presto.spi.plan.PartitioningScheme; import com.facebook.presto.spi.plan.PlanNode; import com.facebook.presto.spi.plan.PlanNodeId; import com.facebook.presto.spi.plan.PlanNodeIdAllocator; import com.facebook.presto.spi.plan.ProjectNode; import com.facebook.presto.spi.plan.SortNode; import com.facebook.presto.spi.plan.TableScanNode; +import com.facebook.presto.spi.plan.TableWriterNode; import com.facebook.presto.spi.plan.ValuesNode; import com.facebook.presto.spi.plan.WindowNode; import com.facebook.presto.spi.relation.CallExpression; @@ -55,34 +65,55 @@ import com.facebook.presto.sql.analyzer.RelationType; import com.facebook.presto.sql.analyzer.Scope; import com.facebook.presto.sql.parser.SqlParser; +import com.facebook.presto.sql.planner.plan.AssignUniqueId; import com.facebook.presto.sql.planner.plan.GroupIdNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.OffsetNode; import com.facebook.presto.sql.planner.plan.UpdateNode; +import com.facebook.presto.sql.tree.BooleanLiteral; import com.facebook.presto.sql.tree.Cast; +import com.facebook.presto.sql.tree.CoalesceExpression; import com.facebook.presto.sql.tree.ComparisonExpression; import com.facebook.presto.sql.tree.Delete; import com.facebook.presto.sql.tree.Expression; import com.facebook.presto.sql.tree.FieldReference; import com.facebook.presto.sql.tree.FrameBound; import com.facebook.presto.sql.tree.FunctionCall; +import com.facebook.presto.sql.tree.GenericLiteral; import com.facebook.presto.sql.tree.GroupingOperation; import com.facebook.presto.sql.tree.IfExpression; import com.facebook.presto.sql.tree.IntervalLiteral; +import com.facebook.presto.sql.tree.IsNotNullPredicate; +import com.facebook.presto.sql.tree.IsNullPredicate; +import com.facebook.presto.sql.tree.Join; import com.facebook.presto.sql.tree.LambdaArgumentDeclaration; import com.facebook.presto.sql.tree.LambdaExpression; +import com.facebook.presto.sql.tree.LogicalBinaryExpression; import com.facebook.presto.sql.tree.LongLiteral; +import com.facebook.presto.sql.tree.Merge; +import com.facebook.presto.sql.tree.MergeCase; +import com.facebook.presto.sql.tree.MergeInsert; +import com.facebook.presto.sql.tree.MergeUpdate; import com.facebook.presto.sql.tree.Node; import com.facebook.presto.sql.tree.NodeLocation; import com.facebook.presto.sql.tree.NodeRef; +import com.facebook.presto.sql.tree.NotExpression; +import com.facebook.presto.sql.tree.NullLiteral; import com.facebook.presto.sql.tree.Offset; import com.facebook.presto.sql.tree.OrderBy; import com.facebook.presto.sql.tree.QualifiedName; import com.facebook.presto.sql.tree.Query; import com.facebook.presto.sql.tree.QuerySpecification; +import com.facebook.presto.sql.tree.Row; +import com.facebook.presto.sql.tree.SearchedCaseExpression; import com.facebook.presto.sql.tree.SortItem; import com.facebook.presto.sql.tree.StringLiteral; +import com.facebook.presto.sql.tree.SubscriptExpression; import com.facebook.presto.sql.tree.SymbolReference; +import com.facebook.presto.sql.tree.Table; import com.facebook.presto.sql.tree.Update; +import com.facebook.presto.sql.tree.WhenClause; import com.facebook.presto.sql.tree.Window; import com.facebook.presto.sql.tree.WindowFrame; import com.google.common.base.VerifyException; @@ -105,8 +136,12 @@ import static com.facebook.presto.SystemSessionProperties.isSkipRedundantSort; import static com.facebook.presto.common.type.BigintType.BIGINT; import static com.facebook.presto.common.type.BooleanType.BOOLEAN; +import static com.facebook.presto.common.type.IntegerType.INTEGER; +import static com.facebook.presto.common.type.TinyintType.TINYINT; import static com.facebook.presto.common.type.VarbinaryType.VARBINARY; import static com.facebook.presto.common.type.VarcharType.VARCHAR; +import static com.facebook.presto.spi.ConnectorMergeSink.INSERT_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.UPDATE_OPERATION_NUMBER; import static com.facebook.presto.spi.StandardErrorCode.INVALID_LIMIT_CLAUSE; import static com.facebook.presto.spi.plan.AggregationNode.groupingSets; import static com.facebook.presto.spi.plan.AggregationNode.singleGroupingSet; @@ -114,11 +149,14 @@ import static com.facebook.presto.spi.plan.ProjectNode.Locality.LOCAL; import static com.facebook.presto.sql.NodeUtils.getSortItemsFromOrderBy; import static com.facebook.presto.sql.analyzer.ExpressionAnalyzer.isNumericType; +import static com.facebook.presto.sql.analyzer.ExpressionTreeUtils.createSymbolReference; import static com.facebook.presto.sql.analyzer.ExpressionTreeUtils.getSourceLocation; +import static com.facebook.presto.sql.planner.PlanBuilder.newPlanBuilder; import static com.facebook.presto.sql.planner.PlannerUtils.newVariable; import static com.facebook.presto.sql.planner.PlannerUtils.toOrderingScheme; import static com.facebook.presto.sql.planner.PlannerUtils.toSortOrder; import static com.facebook.presto.sql.planner.PlannerUtils.toVariableReference; +import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION; import static com.facebook.presto.sql.planner.TranslateExpressionsUtil.analyzeCallExpressionTypes; import static com.facebook.presto.sql.planner.TranslateExpressionsUtil.toRowExpression; import static com.facebook.presto.sql.planner.optimizations.WindowNodeUtil.toBoundType; @@ -395,6 +433,430 @@ public UpdateNode plan(Update node) outputs); } + /** + * Plan a MERGE statement. The MERGE statement is processed by creating a RIGHT JOIN between the target table and the source. + * Example of converting a MERGE statement into a SELECT statement with a RIGHT JOIN: + * Merge statement: + * MERGE INTO t USING s + * ON (t. = s.) + * WHEN MATCHED THEN + * UPDATE SET = s. + t., + * = s. + * WHEN NOT MATCHED THEN + * INSERT (column1, column2, column3) + * VALUES (s.column1, s.column2, s.column3); + * + * SELECT statement with a RIGHT JOIN created to process the previous MERGE statement: + * SELECT + * CASE + * WHEN present THEN + * -- Update column values: present=true, operation UPDATE=3, case_number=0 + * row(t.column1, s.column1 + t.column1, s.column2, true, 3, 0) + * WHEN (present IS NULL) THEN + * -- Insert column values: present=false; operation INSERT=1, case_number=1 + * row(s.column1, s.column2, s.column3, false, 1, 1) + * ELSE + * -- Null values for no case matched: present=false, operation=-1, case_number=-1 + * row(null, null, null, false, -1, -1) + * END + * FROM + * (SELECT *, true AS present FROM ) t + * RIGHT JOIN s + * ON s. = t.; + * + * @param mergeStmt the MERGE statement to plan into a MergeWriterNode. + * @return a MergeWriterNode that represents the plan for the MERGE statement. + */ + public MergeWriterNode plan(Merge mergeStmt) + { + // The goal of this method is to build the following MERGE INTO execution plan: + // + // MergeWriterNode : Write the merge results into the target table. + // | + // MergeProcessorNode : Processes the result of the RIGHT JOIN to identify which rows need to be inserted and which need to be updated. + // | + // FilterDuplicateMatchingRows : Look for marked rows in the previous step. If it finds one, then it stops MERGE execution and returns an error. + // | + // MarkDistinctNode : Look for target rows that matched more than one source row and mark them. + // | + // RightEquiJoin : Run a RIGHT JOIN as a first step to process the MERGE INTO command. + // / \ + // ProjectPresentColumn \ : Add a constant "TRUE" column to the target table. It is used to identify matching rows. + // | \ + // AssignUniqueID \ : Assign a unique ID to each row in the target table. + // | \ + // TableScan TableScan : Read data from the target and source tables. + // | | + // (target table) (source table) + + Analysis.MergeAnalysis mergeAnalysis = analysis.getMergeAnalysis().orElseThrow(() -> new IllegalArgumentException("analysis.getMergeAnalysis() isn't present")); + + // Make the plan for the merge target table scan + RelationPlan targetTableRelationPlan = new RelationPlanner(analysis, variableAllocator, idAllocator, lambdaDeclarationToVariableMap, metadata, session, sqlParser) + .process(mergeStmt.getTarget(), sqlPlannerContext); + + // Assign a unique id to every target table row + VariableReferenceExpression uniqueIdVariable = variableAllocator.newVariable("unique_id", BIGINT); + AssignUniqueId assignUniqueRowIdToTargetTable = new AssignUniqueId(getSourceLocation(mergeStmt), idAllocator.getNextId(), targetTableRelationPlan.getRoot(), uniqueIdVariable); + RelationPlan relationPlanWithUniqueRowIds = new RelationPlan( + assignUniqueRowIdToTargetTable, + mergeAnalysis.getTargetTableScope(), + targetTableRelationPlan.getFieldMappings()); + + // TODO #20578: Do we really need the "present" column? Can we use the "unique_id" column to verify if there is a row match between the target and source table? + // Project the "present" column + Assignments.Builder targetTableProjections = Assignments.builder(); + relationPlanWithUniqueRowIds.getRoot().getOutputVariables().forEach(variable -> targetTableProjections.put(variable, variable)); + + VariableReferenceExpression presentColumn = variableAllocator.newVariable("present", BOOLEAN); + targetTableProjections.put(presentColumn, rowExpression(TRUE_LITERAL, sqlPlannerContext)); + SymbolReference presentColumnSymbolReference = createSymbolReference(presentColumn); + + ProjectNode projectNodeWithUniqueRowIdsAndPresentColumn = + new ProjectNode(idAllocator.getNextId(), relationPlanWithUniqueRowIds.getRoot(), targetTableProjections.build()); + + RelationPlan planWithWithUniqueRowIdsAndPresentColumn = new RelationPlan( + projectNodeWithUniqueRowIdsAndPresentColumn, + mergeAnalysis.getTargetTableScope(), + relationPlanWithUniqueRowIds.getFieldMappings()); + + RelationPlan sourceRelationPlan = new RelationPlanner(analysis, variableAllocator, idAllocator, lambdaDeclarationToVariableMap, metadata, session, sqlParser) + .process(mergeStmt.getSource(), sqlPlannerContext); + + RelationPlan joinRelationPlan = new RelationPlanner(analysis, variableAllocator, idAllocator, lambdaDeclarationToVariableMap, metadata, session, sqlParser) + .planJoin( + coerceIfNecessary(analysis, mergeStmt.getPredicate(), mergeStmt.getPredicate()), + Join.Type.RIGHT, mergeAnalysis.getJoinScope(), + planWithWithUniqueRowIdsAndPresentColumn, sourceRelationPlan, + mergeStmt, sqlPlannerContext); + + // Build the SearchedCaseExpression that creates the project "merge_row" + PlanBuilder joinSubPlan = newPlanBuilder(joinRelationPlan, analysis, lambdaDeclarationToVariableMap); + + // CASE + // WHEN present THEN row(column1, column2, ..., present=true, operation UPDATE=3, case_number=0) + // WHEN (present IS NULL) THEN row(column1, column2, ..., present=false, operation INSERT=1, case_number=1) + // ELSE row(null, null, ..., false, -1, -1) + // END + ImmutableList.Builder whenClauses = ImmutableList.builder(); + for (int caseNumber = 0; caseNumber < mergeStmt.getMergeCases().size(); caseNumber++) { + MergeCase mergeCase = mergeStmt.getMergeCases().get(caseNumber); + + ImmutableList.Builder joinResultBuilder = ImmutableList.builder(); + List> mergeCaseColumnsHandles = mergeAnalysis.getMergeCaseColumnHandles(); + List mergeCaseSetColumns = mergeCaseColumnsHandles.get(caseNumber); + for (ColumnHandle targetColumnHandle : mergeAnalysis.getTargetColumnHandles()) { + int index = mergeCaseSetColumns.indexOf(targetColumnHandle); + int fieldNumber = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(targetColumnHandle), "Field number for ColumnHandle is null"); + Expression expression; + if (index >= 0) { // Update column value + Expression setExpression = mergeCase.getSetExpressions().get(index); + joinSubPlan = subqueryPlanner.handleSubqueries(joinSubPlan, setExpression, mergeStmt, sqlPlannerContext); + expression = joinSubPlan.rewrite(setExpression); + expression = coerceIfNecessary(analysis, setExpression, expression); + expression = checkNotNullColumns(targetColumnHandle, expression, fieldNumber, mergeAnalysis); + } + else { // Insert column value + expression = createSymbolReference(planWithWithUniqueRowIdsAndPresentColumn.getFieldMappings().get(fieldNumber)); + + if (mergeCase instanceof MergeInsert) { + expression = checkNotNullColumns(targetColumnHandle, expression, fieldNumber, mergeAnalysis); + } + } + joinResultBuilder.add(expression); + } + + // Add the "present" column value. It is a boolean column which is true if a target table row was matched. + joinResultBuilder.add(new IsNotNullPredicate(presentColumnSymbolReference)); + + // Add the operation number + joinResultBuilder.add(new GenericLiteral("TINYINT", String.valueOf(getMergeCaseOperationNumber(mergeCase)))); + + // Add the mergeStmt case number, needed by MarkDistinct + joinResultBuilder.add(new GenericLiteral("INTEGER", String.valueOf(caseNumber))); + + // Build the match condition for the MERGE case + Expression mergeCondition = mergeCase instanceof MergeInsert ? + new IsNullPredicate(presentColumnSymbolReference) : presentColumnSymbolReference; + + whenClauses.add(new WhenClause(mergeCondition, new Row(joinResultBuilder.build()))); + } + + // Build the "else" clause for the SearchedCaseExpression + ImmutableList.Builder joinElseBuilder = ImmutableList.builder(); + mergeAnalysis.getTargetColumnsMetadata().forEach(columnMetadata -> + joinElseBuilder.add(new Cast(new NullLiteral(), columnMetadata.getType().getDisplayName()))); + + // Add the "present" column value. It is always FALSE for the "else" clause. + joinElseBuilder.add(BooleanLiteral.FALSE_LITERAL); + // The operation number column value: -1 + joinElseBuilder.add(new GenericLiteral("TINYINT", "-1")); + // The case number column value: -1 + joinElseBuilder.add(new GenericLiteral("INTEGER", "-1")); + + SearchedCaseExpression caseExpression = new SearchedCaseExpression(whenClauses.build(), Optional.of(new Row(joinElseBuilder.build()))); + + RowType mergeRowType = createMergeRowType(mergeAnalysis.getTargetColumnsMetadata()); + Table targetTable = mergeAnalysis.getTargetTable(); + FieldReference rowIdReference = analysis.getRowIdField(targetTable); + + VariableReferenceExpression mergeRowIdVariable = planWithWithUniqueRowIdsAndPresentColumn.getFieldMappings().get(rowIdReference.getFieldIndex()); + VariableReferenceExpression mergeRowVariable = variableAllocator.newVariable("merge_row", mergeRowType); + + // Project the partitioning variables, the merge_row, the rowId, and the unique_id variable. + Assignments.Builder projectionAssignmentsBuilder = Assignments.builder(); + for (ColumnHandle column : mergeAnalysis.getTargetRedistributionColumnHandles()) { + int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(column), "Could not find fieldIndex for redistribution column"); + VariableReferenceExpression variable = planWithWithUniqueRowIdsAndPresentColumn.getFieldMappings().get(fieldIndex); + projectionAssignmentsBuilder.put(variable, variable); + } + + projectionAssignmentsBuilder.put(uniqueIdVariable, uniqueIdVariable); + projectionAssignmentsBuilder.put(mergeRowIdVariable, mergeRowIdVariable); + projectionAssignmentsBuilder.put(mergeRowVariable, rowExpression(caseExpression, sqlPlannerContext)); + + ProjectNode joinSubPlanProject = new ProjectNode( + idAllocator.getNextId(), + joinSubPlan.getRoot(), + projectionAssignmentsBuilder.build()); + + // Now add a column for the case_number, gotten from the merge_row + SubscriptExpression caseNumberExpression = new SubscriptExpression( + createSymbolReference(mergeRowVariable), new LongLiteral(Long.toString(mergeRowType.getFields().size()))); + + VariableReferenceExpression caseNumberVariable = variableAllocator.newVariable("case_number", INTEGER); + + ProjectNode joinProjectNode = new ProjectNode( + joinSubPlanProject.getSourceLocation(), + idAllocator.getNextId(), + joinSubPlanProject, + Assignments.builder() + .putAll(joinSubPlanProject.getOutputVariables().stream().collect(toImmutableMap(Function.identity(), Function.identity()))) + .put(caseNumberVariable, rowExpression(caseNumberExpression, sqlPlannerContext)) + .build(), + LOCAL); // TODO #20578: Is LOCAL the correct value? + + // Mark distinct combinations of the unique_id value and the case_number + VariableReferenceExpression isDistinctVariable = variableAllocator.newVariable("is_distinct", BOOLEAN); + MarkDistinctNode markDistinctNode = new MarkDistinctNode( + getSourceLocation(mergeStmt), idAllocator.getNextId(), joinProjectNode, isDistinctVariable, + ImmutableList.of(uniqueIdVariable, caseNumberVariable), Optional.empty()); + + // Raise an error if unique_id variable is non-null and the unique_id/case_number combination was not distinct + Expression multipleMatchesExpression = new IfExpression( + LogicalBinaryExpression.and( + new NotExpression(createSymbolReference(isDistinctVariable)), + new IsNotNullPredicate(createSymbolReference(uniqueIdVariable))), + new Cast( + new FunctionCall( + QualifiedName.of("presto", "default", "fail"), + ImmutableList.of(new Cast(new StringLiteral( + "MERGE INTO operation failed for target table '" + targetTable.getName() + "'. " + + "One or more rows in the target table matched multiple source rows. " + + "The MERGE INTO command requires each target row to match at most one source row. " + + "Please review the ON condition to ensure it produces a one-to-one or one-to-none match."), + VARCHAR.getTypeSignature().toString()))), + BOOLEAN.getTypeSignature().toString()), + TRUE_LITERAL); + + FilterNode filterMultipleMatches = new FilterNode(getSourceLocation(mergeStmt), idAllocator.getNextId(), + markDistinctNode, rowExpression(multipleMatchesExpression, sqlPlannerContext)); + + TableHandle targetTableHandle = analysis.getTableHandle(targetTable); + RowChangeParadigm rowChangeParadigm = metadata.getRowChangeParadigm(session, targetTableHandle); + Type rowIdType = analysis.getType(analysis.getRowIdField(targetTable)); + TableMetadata targetTableMetadata = metadata.getTableMetadata(session, targetTableHandle); + + List targetColumnsDataTypes = targetTableMetadata.getMetadata().getColumns().stream() + .filter(column -> !column.isHidden()) + .map(ColumnMetadata::getType) + .collect(toImmutableList()); + + TableWriterNode.MergeParadigmAndTypes mergeParadigmAndTypes = + new TableWriterNode.MergeParadigmAndTypes(rowChangeParadigm, targetColumnsDataTypes, rowIdType); + + Optional mergeHandle = Optional.of(metadata.beginMerge(session, targetTableHandle)); + TableWriterNode.MergeTarget mergeTarget = + new TableWriterNode.MergeTarget(targetTableHandle, mergeHandle, targetTableMetadata.getTable(), mergeParadigmAndTypes); + + ImmutableList.Builder mergeColumnVariablesBuilder = ImmutableList.builder(); + for (ColumnHandle columnHandle : mergeAnalysis.getTargetColumnHandles()) { + int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(columnHandle), "Could not find field number for column handle"); + mergeColumnVariablesBuilder.add(planWithWithUniqueRowIdsAndPresentColumn.getFieldMappings().get(fieldIndex)); + } + List mergeColumnVariables = mergeColumnVariablesBuilder.build(); + + ImmutableList.Builder mergeRedistributionVariablesBuilder = ImmutableList.builder(); + for (ColumnHandle columnHandle : mergeAnalysis.getTargetRedistributionColumnHandles()) { + int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(columnHandle), "Could not find field number for column handle"); + mergeRedistributionVariablesBuilder.add(planWithWithUniqueRowIdsAndPresentColumn.getFieldMappings().get(fieldIndex)); + } + + // Variable to specify whether the MERGE INTO statement should insert a new row or update an existing one. + // Operations defined in ConnectorMergeSink: INSERT_OPERATION_NUMBER and UPDATE_OPERATION_NUMBER. + VariableReferenceExpression mergeOperationVariable = variableAllocator.newVariable("operation", TINYINT); + + // Variable to indicate whether the row is an insert resulting from an UPDATE or an INSERT. + // The RowChangeParadigm.DELETE_ROW_AND_INSERT_ROW will use this information. + // Values: + // 1: if this row is an insert derived from an UPDATE + // 0: if this row is an insert derived from an INSERT + VariableReferenceExpression insertFromUpdateVariable = variableAllocator.newVariable("insert_from_update", TINYINT); + + List mergeProjectedVariables = ImmutableList.builder() + .addAll(mergeColumnVariables) + .add(mergeOperationVariable) + .add(mergeRowIdVariable) + .add(insertFromUpdateVariable) + .build(); + + MergeProcessorNode mergeProcessorNode = new MergeProcessorNode( + getSourceLocation(mergeStmt), + idAllocator.getNextId(), + filterMultipleMatches, + mergeTarget, + mergeRowIdVariable, + mergeRowVariable, + mergeColumnVariables, + mergeRedistributionVariablesBuilder.build(), + mergeProjectedVariables); + + Optional mergeWriterPartitioningScheme = createMergePartitioningScheme( + mergeAnalysis.getInsertLayout(), + mergeColumnVariables, + mergeAnalysis.getInsertPartitioningArgumentIndexes(), + mergeAnalysis.getUpdateLayout(), + mergeRowIdVariable, + mergeOperationVariable); + + List mergeWriterOutputs = ImmutableList.of( + variableAllocator.newVariable("partialrows", BIGINT), + variableAllocator.newVariable("fragment", VARBINARY)); + + return new MergeWriterNode( + getSourceLocation(mergeStmt), + idAllocator.getNextId(), + mergeProcessorNode, + mergeTarget, + mergeProjectedVariables, + mergeWriterPartitioningScheme, + mergeWriterOutputs); + } + + /** + * Method to create a CoalesceExpression that triggers an error if the query attempts to insert a NULL value + * into a non-nullable column. The same applies if the query tries to update a non-nullable column with a NULL value. + * + * @return The same expression if the column allows null values; otherwise, a CoalesceExpression that + * prevents inserting NULL values into non-nullable columns. + */ + private static Expression checkNotNullColumns( + ColumnHandle targetColumnHandle, + Expression expression, + int fieldNumber, + Analysis.MergeAnalysis mergeAnalysis) + { + // If the current column allows NULL values, then the method returns the original expression. + if (!mergeAnalysis.getNonNullableColumnHandles().contains(targetColumnHandle)) { + return expression; + } + + ColumnMetadata columnMetadata = mergeAnalysis.getTargetColumnsMetadata().get(fieldNumber); + + // Build a coalesce expression that returns an error when the original expression value is NULL. + return new CoalesceExpression(expression, + new Cast( + new FunctionCall( + QualifiedName.of("presto", "default", "fail"), + ImmutableList.of(new Cast(new StringLiteral( + "NULL value not allowed for NOT NULL column. Table: " + mergeAnalysis.getTargetTable().getName() + + " Column: " + columnMetadata.getName()), + VARCHAR.getTypeSignature().toString()))), + columnMetadata.getType().getTypeSignature().toString())); + } + + private static int getMergeCaseOperationNumber(MergeCase mergeCase) + { + if (mergeCase instanceof MergeInsert) { + return INSERT_OPERATION_NUMBER; + } + if (mergeCase instanceof MergeUpdate) { + return UPDATE_OPERATION_NUMBER; + } + throw new IllegalArgumentException("Unrecognized MergeCase: " + mergeCase); + } + + private static RowType createMergeRowType(List allColumnsMetadata) + { + // Create the RowType that holds all column values + List fields = new ArrayList<>(); + for (ColumnMetadata columnMetadata : allColumnsMetadata) { + fields.add(new RowType.Field(Optional.empty(), columnMetadata.getType())); + } + fields.add(new RowType.Field(Optional.empty(), BOOLEAN)); // present + fields.add(new RowType.Field(Optional.empty(), TINYINT)); // operation_number + fields.add(new RowType.Field(Optional.empty(), INTEGER)); // case_number + return RowType.from(fields); + } + + public static Optional createMergePartitioningScheme( + Optional insertLayout, + List variables, + List insertPartitioningArgumentIndexes, + Optional updateLayout, + VariableReferenceExpression rowIdVariable, + VariableReferenceExpression operationVariable) + { + if (!insertLayout.isPresent() && !updateLayout.isPresent()) { + return Optional.empty(); + } + + Optional insertPartitioning = insertLayout.map(layout -> { + List arguments = insertPartitioningArgumentIndexes.stream() + .map(variables::get) + .collect(toImmutableList()); + + return layout.getPartitioning() + .map(handle -> new PartitioningScheme(Partitioning.create(handle, arguments), variables)) + // empty connector partitioning handle means evenly partitioning on partitioning columns + .orElseGet(() -> new PartitioningScheme(Partitioning.create(FIXED_HASH_DISTRIBUTION, arguments), variables)); + }); + + Optional updatePartitioning = updateLayout.map(handle -> + new PartitioningScheme(Partitioning.create(handle, ImmutableList.of(rowIdVariable)), ImmutableList.of(rowIdVariable))); + + PartitioningHandle partitioningHandle = new PartitioningHandle( + Optional.empty(), + Optional.empty(), + new MergePartitioningHandle(insertPartitioning, updatePartitioning)); + + List combinedVariables = new ArrayList<>(); + combinedVariables.add(operationVariable); + insertPartitioning.ifPresent(scheme -> combinedVariables.addAll(partitioningVariables(scheme))); + updatePartitioning.ifPresent(scheme -> combinedVariables.addAll(partitioningVariables(scheme))); + + return Optional.of(new PartitioningScheme(Partitioning.create(partitioningHandle, combinedVariables), combinedVariables)); + } + + private static List partitioningVariables(PartitioningScheme scheme) + { + return new ArrayList<>(scheme.getPartitioning().getVariableReferences()); + } + + public static Expression coerceIfNecessary(Analysis analysis, Expression original, Expression rewritten) + { + Type coercion = analysis.getCoercion(original); + if (coercion == null) { + return rewritten; + } + + return new Cast( + rewritten, + coercion.getDisplayName(), + false, + analysis.isTypeOnlyCoercion(original)); + } + private Optional getIdForLeftTableScan(PlanNode node) { if (node instanceof TableScanNode) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java index 0bd4e47fc67b8..636e6195b8cc4 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java @@ -219,7 +219,8 @@ protected RelationPlan visitTable(Table node, SqlPlannerContext context) ImmutableList.Builder outputVariablesBuilder = ImmutableList.builder(); ImmutableMap.Builder columns = ImmutableMap.builder(); for (Field field : scope.getRelationType().getAllFields()) { - VariableReferenceExpression variable = variableAllocator.newVariable(getSourceLocation(node), field.getName().get(), field.getType()); + // TODO #20578: Check the consequences of adding .orElse("field") in field.getName(). + VariableReferenceExpression variable = variableAllocator.newVariable(getSourceLocation(node), field.getName().orElse("field"), field.getType()); outputVariablesBuilder.add(variable); columns.put(variable, analysis.getColumn(field)); } @@ -423,6 +424,11 @@ protected RelationPlan visitJoin(Join node, SqlPlannerContext context) return planJoinUsing(node, leftPlan, rightPlan, context); } + return planJoin(analysis.getJoinCriteria(node), node.getType(), analysis.getScope(node), leftPlan, rightPlan, node, context); + } + + public RelationPlan planJoin(Expression criteria, Join.Type type, Scope scope, RelationPlan leftPlan, RelationPlan rightPlan, Node node, SqlPlannerContext context) + { PlanBuilder leftPlanBuilder = initializePlanBuilder(leftPlan); PlanBuilder rightPlanBuilder = initializePlanBuilder(rightPlan); @@ -436,12 +442,10 @@ protected RelationPlan visitJoin(Join node, SqlPlannerContext context) List complexJoinExpressions = new ArrayList<>(); List postInnerJoinConditions = new ArrayList<>(); - if (node.getType() != Join.Type.CROSS && node.getType() != Join.Type.IMPLICIT) { - Expression criteria = analysis.getJoinCriteria(node); - - RelationType left = analysis.getOutputDescriptor(node.getLeft()); - RelationType right = analysis.getOutputDescriptor(node.getRight()); + RelationType left = leftPlan.getDescriptor(); + RelationType right = rightPlan.getDescriptor(); + if (type != Join.Type.CROSS && type != Join.Type.IMPLICIT) { List leftComparisonExpressions = new ArrayList<>(); List rightComparisonExpressions = new ArrayList<>(); List joinConditionComparisonOperators = new ArrayList<>(); @@ -449,7 +453,7 @@ protected RelationPlan visitJoin(Join node, SqlPlannerContext context) for (Expression conjunct : ExpressionUtils.extractConjuncts(criteria)) { conjunct = ExpressionUtils.normalize(conjunct); - if (!isEqualComparisonExpression(conjunct) && node.getType() != INNER) { + if (!isEqualComparisonExpression(conjunct) && type != INNER) { complexJoinExpressions.add(conjunct); continue; } @@ -513,7 +517,7 @@ else if (firstDependencies.stream().allMatch(right::canResolve) && secondDepende PlanNode root = new JoinNode( getSourceLocation(node), idAllocator.getNextId(), - JoinNodeUtils.typeConvert(node.getType()), + JoinNodeUtils.typeConvert(type), leftPlanBuilder.getRoot(), rightPlanBuilder.getRoot(), equiClauses.build(), @@ -527,7 +531,7 @@ else if (firstDependencies.stream().allMatch(right::canResolve) && secondDepende Optional.empty(), ImmutableMap.of()); - if (node.getType() != INNER) { + if (type != INNER) { for (Expression complexExpression : complexJoinExpressions) { Set inPredicates = subqueryPlanner.collectInPredicateSubqueries(complexExpression, node); if (!inPredicates.isEmpty()) { @@ -536,10 +540,7 @@ else if (firstDependencies.stream().allMatch(right::canResolve) && secondDepende } } - if (node.getType() == LEFT || node.getType() == RIGHT) { - RelationType left = analysis.getOutputDescriptor(node.getLeft()); - RelationType right = analysis.getOutputDescriptor(node.getRight()); - + if (type == LEFT || type == RIGHT) { for (Expression complexJoinExpression : complexJoinExpressions) { Set dependencies = VariablesExtractor.extractNames(complexJoinExpression, analysis.getColumnReferences()); // If there are no dependencies, no subqueries, or if the expression references both inputs, @@ -555,7 +556,7 @@ else if (firstDependencies.stream().allMatch(right::canResolve) && secondDepende // If the subquery references the right input, those variables will remain unresolved and caught in NoIdentifierLeftChecker leftPlanBuilder = subqueryPlanner.handleUncorrelatedSubqueries(leftPlanBuilder, ImmutableList.of(complexJoinExpression), node, context); } - else if (node.getType() == LEFT && !dependencies.stream().allMatch(left::canResolve)) { + else if (type == LEFT && !dependencies.stream().allMatch(left::canResolve)) { rightPlanBuilder = subqueryPlanner.handleSubqueries(rightPlanBuilder, complexJoinExpression, node, context); } else { @@ -570,19 +571,19 @@ else if (node.getType() == LEFT && !dependencies.stream().allMatch(left::canReso } } - RelationPlan intermediateRootRelationPlan = new RelationPlan(root, analysis.getScope(node), outputs); + RelationPlan intermediateRootRelationPlan = new RelationPlan(root, scope, outputs); TranslationMap translationMap = new TranslationMap(intermediateRootRelationPlan, analysis, lambdaDeclarationToVariableMap); translationMap.setFieldMappings(outputs); translationMap.putExpressionMappingsFrom(leftPlanBuilder.getTranslations()); translationMap.putExpressionMappingsFrom(rightPlanBuilder.getTranslations()); - if (node.getType() != INNER && !complexJoinExpressions.isEmpty()) { + if (type != INNER && !complexJoinExpressions.isEmpty()) { Expression joinedFilterCondition = ExpressionUtils.and(complexJoinExpressions); Expression rewrittenFilterCondition = translationMap.rewrite(joinedFilterCondition); root = new JoinNode( getSourceLocation(node), idAllocator.getNextId(), - JoinNodeUtils.typeConvert(node.getType()), + JoinNodeUtils.typeConvert(type), leftPlanBuilder.getRoot(), rightPlanBuilder.getRoot(), equiClauses.build(), @@ -597,7 +598,7 @@ else if (node.getType() == LEFT && !dependencies.stream().allMatch(left::canReso ImmutableMap.of()); } - if (node.getType() == INNER) { + if (type == INNER) { // rewrite all the other conditions using output variables from left + right plan node. PlanBuilder rootPlanBuilder = new PlanBuilder(translationMap, root); rootPlanBuilder = subqueryPlanner.handleSubqueries(rootPlanBuilder, complexJoinExpressions, node, context); @@ -614,7 +615,7 @@ else if (node.getType() == LEFT && !dependencies.stream().allMatch(left::canReso } } - return new RelationPlan(root, analysis.getScope(node), outputs); + return new RelationPlan(root, scope, outputs); } private RelationPlan planJoinUsing(Join node, RelationPlan left, RelationPlan right, SqlPlannerContext context) diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java index 03522fabd5303..f0ca7b8c1cac1 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java @@ -54,6 +54,8 @@ import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; @@ -410,6 +412,18 @@ public Map visitPlan(PlanNode node, Context context) { throw new UnsupportedOperationException("not yet implemented: " + node.getClass().getName()); } + + @Override + public Map visitMergeWriter(MergeWriterNode node, Context context) + { + return node.getSource().accept(this, context); + } + + @Override + public Map visitMergeProcessor(MergeProcessorNode node, Context context) + { + return node.getSource().accept(this, context); + } } private static class Context diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SystemPartitioningHandle.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SystemPartitioningHandle.java index 8a4174c9cefe7..d7b859ec00348 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SystemPartitioningHandle.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/SystemPartitioningHandle.java @@ -13,43 +13,31 @@ */ package com.facebook.presto.sql.planner; -import com.facebook.presto.Session; import com.facebook.presto.common.Page; import com.facebook.presto.common.type.Type; -import com.facebook.presto.execution.scheduler.NodeScheduler; -import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelector; -import com.facebook.presto.metadata.InternalNode; import com.facebook.presto.operator.BucketPartitionFunction; import com.facebook.presto.operator.HashGenerator; import com.facebook.presto.operator.InterpretedHashGenerator; import com.facebook.presto.operator.PartitionFunction; import com.facebook.presto.operator.PrecomputedHashGenerator; import com.facebook.presto.spi.BucketFunction; -import com.facebook.presto.spi.Node; import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; import com.facebook.presto.spi.plan.PartitioningHandle; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.common.collect.ImmutableList; import java.util.List; import java.util.Objects; import java.util.Optional; -import java.util.function.Predicate; -import static com.facebook.presto.SystemSessionProperties.getHashPartitionCount; -import static com.facebook.presto.SystemSessionProperties.getMaxTasksPerStage; -import static com.facebook.presto.spi.StandardErrorCode.NO_NODES_AVAILABLE; -import static com.facebook.presto.util.Failures.checkCondition; import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; -import static java.lang.Math.min; import static java.util.Objects.requireNonNull; public final class SystemPartitioningHandle implements ConnectorPartitioningHandle { - private enum SystemPartitioning + enum SystemPartitioning { SINGLE, FIXED, @@ -151,30 +139,6 @@ public String toString() return partitioning.toString(); } - public NodePartitionMap getNodePartitionMap(Session session, NodeScheduler nodeScheduler, Optional> nodePredicate) - { - NodeSelector nodeSelector = nodeScheduler.createNodeSelector(session, null, nodePredicate); - List nodes; - if (partitioning == SystemPartitioning.COORDINATOR_ONLY) { - nodes = ImmutableList.of(nodeSelector.selectCurrentNode()); - } - else if (partitioning == SystemPartitioning.SINGLE) { - nodes = nodeSelector.selectRandomNodes(1); - } - else if (partitioning == SystemPartitioning.FIXED) { - nodes = nodeSelector.selectRandomNodes(min(getHashPartitionCount(session), getMaxTasksPerStage(session))); - } - else { - throw new IllegalArgumentException("Unsupported plan distribution " + partitioning); - } - - checkCondition(!nodes.isEmpty(), NO_NODES_AVAILABLE, "No worker nodes available"); - - return new NodePartitionMap(nodes, split -> { - throw new UnsupportedOperationException("System distribution does not support source splits"); - }); - } - public PartitionFunction getPartitionFunction(List partitionChannelTypes, boolean isHashPrecomputed, int[] bucketToPartition) { requireNonNull(partitionChannelTypes, "partitionChannelTypes is null"); @@ -252,7 +216,7 @@ public int getBucket(Page page, int position) } } - private static class RoundRobinBucketFunction + public static class RoundRobinBucketFunction implements BucketFunction { private final int bucketCount; diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/PruneMergeSourceColumns.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/PruneMergeSourceColumns.java new file mode 100644 index 0000000000000..9237e98e06833 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/PruneMergeSourceColumns.java @@ -0,0 +1,43 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner.iterative.rule; + +import com.facebook.presto.matching.Captures; +import com.facebook.presto.matching.Pattern; +import com.facebook.presto.sql.planner.iterative.Rule; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; +import com.google.common.collect.ImmutableSet; + +import static com.facebook.presto.sql.planner.iterative.rule.Util.restrictChildOutputs; +import static com.facebook.presto.sql.planner.plan.Patterns.mergeWriter; + +public class PruneMergeSourceColumns + implements Rule +{ + private static final Pattern PATTERN = mergeWriter(); + + @Override + public Pattern getPattern() + { + return PATTERN; + } + + @Override + public Result apply(MergeWriterNode mergeNode, Captures captures, Context context) + { + return restrictChildOutputs(context.getIdAllocator(), mergeNode, ImmutableSet.copyOf(mergeNode.getMergeProcessorProjectedVariables())) + .map(Result::ofPlanNode) + .orElse(Result.empty()); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java index 61b53b70905f0..80a80a7f9d114 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java @@ -74,6 +74,7 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; @@ -753,8 +754,18 @@ public PlanWithProperties visitTableWriter(TableWriterNode node, PreferredProper { PlanWithProperties source = accept(node.getSource(), preferredProperties); - Optional shufflePartitioningScheme = node.getTablePartitioningScheme(); - if (!node.isSingleWriterPerPartitionRequired()) { + source = getWriterPlanWithProperties(node.getTablePartitioningScheme(), source, node.isSingleWriterPerPartitionRequired()); + + return rebaseAndDeriveProperties(node, source); + } + + private PlanWithProperties getWriterPlanWithProperties( + Optional nodeTablePartitioningScheme, + PlanWithProperties source, + boolean isSingleWriterPerPartitionRequired) + { + Optional shufflePartitioningScheme = nodeTablePartitioningScheme; + if (!isSingleWriterPerPartitionRequired) { // prefer scale writers if single writer per partition is not required // TODO: take into account partitioning scheme in scale writer tasks implementation if (scaleWriters) { @@ -764,19 +775,20 @@ else if (redistributeWrites) { shufflePartitioningScheme = Optional.of(new PartitioningScheme(Partitioning.create(FIXED_ARBITRARY_DISTRIBUTION, ImmutableList.of()), source.getNode().getOutputVariables())); } else { - return rebaseAndDeriveProperties(node, source); + return source; } } + PlanWithProperties newSource = source; if (shufflePartitioningScheme.isPresent() && // TODO: Deprecate compatible table partitioning - !source.getProperties().isCompatibleTablePartitioningWith(shufflePartitioningScheme.get().getPartitioning(), false, metadata, session) && - !(source.getProperties().isRefinedPartitioningOver(shufflePartitioningScheme.get().getPartitioning(), false, metadata, session) && - canPushdownPartialMerge(source.getNode(), partialMergePushdownStrategy))) { + !newSource.getProperties().isCompatibleTablePartitioningWith(shufflePartitioningScheme.get().getPartitioning(), false, metadata, session) && + !(newSource.getProperties().isRefinedPartitioningOver(shufflePartitioningScheme.get().getPartitioning(), false, metadata, session) && + canPushdownPartialMerge(newSource.getNode(), partialMergePushdownStrategy))) { PartitioningScheme exchangePartitioningScheme = shufflePartitioningScheme.get(); - if (node.getTablePartitioningScheme().isPresent() && isPrestoSparkAssignBucketToPartitionForPartitionedTableWriteEnabled(session)) { + if (nodeTablePartitioningScheme.isPresent() && isPrestoSparkAssignBucketToPartitionForPartitionedTableWriteEnabled(session)) { int writerThreadsPerNode = getTaskPartitionedWriterCount(session); - int bucketCount = getBucketCount(node.getTablePartitioningScheme().get().getPartitioning().getHandle()); + int bucketCount = getBucketCount(nodeTablePartitioningScheme.get().getPartitioning().getHandle()); int[] bucketToPartition = new int[bucketCount]; for (int i = 0; i < bucketCount; i++) { bucketToPartition[i] = i / writerThreadsPerNode; @@ -784,15 +796,27 @@ else if (redistributeWrites) { exchangePartitioningScheme = exchangePartitioningScheme.withBucketToPartition(Optional.of(bucketToPartition)); } - source = withDerivedProperties( + newSource = withDerivedProperties( partitionedExchange( idAllocator.getNextId(), REMOTE_STREAMING, - source.getNode(), + newSource.getNode(), exchangePartitioningScheme), - source.getProperties()); + newSource.getProperties()); } - return rebaseAndDeriveProperties(node, source); + return newSource; + } + + @Override + public PlanWithProperties visitMergeWriter(MergeWriterNode node, PreferredProperties preferredProperties) + { + PlanWithProperties source = node.getSource().accept(this, preferredProperties); + + Optional partitioningScheme = node.getPartitioningScheme(); + boolean isSingleWriterPerPartitionRequired = partitioningScheme.isPresent() && !partitioningScheme.get().isScaleWriters(); + PlanWithProperties partitionedSource = getWriterPlanWithProperties(partitioningScheme, source, isSingleWriterPerPartitionRequired); + + return rebaseAndDeriveProperties(node, partitionedSource); } private int getBucketCount(PartitioningHandle partitioning) diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java index c76c6b6252771..bb86e5bcf5dbd 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java @@ -46,7 +46,9 @@ import com.facebook.presto.spi.plan.TopNNode; import com.facebook.presto.spi.plan.UnionNode; import com.facebook.presto.spi.plan.WindowNode; +import com.facebook.presto.spi.relation.ConstantExpression; import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.facebook.presto.sql.planner.SystemPartitioningHandle; import com.facebook.presto.sql.planner.TypeProvider; import com.facebook.presto.sql.planner.optimizations.StreamPropertyDerivations.StreamProperties; import com.facebook.presto.sql.planner.plan.ApplyNode; @@ -55,6 +57,7 @@ import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; @@ -92,6 +95,7 @@ import static com.facebook.presto.sql.planner.optimizations.StreamPreferredProperties.defaultParallelism; import static com.facebook.presto.sql.planner.optimizations.StreamPreferredProperties.exactlyPartitionedOn; import static com.facebook.presto.sql.planner.optimizations.StreamPreferredProperties.fixedParallelism; +import static com.facebook.presto.sql.planner.optimizations.StreamPreferredProperties.partitionedOn; import static com.facebook.presto.sql.planner.optimizations.StreamPreferredProperties.singleStream; import static com.facebook.presto.sql.planner.optimizations.StreamPropertyDerivations.StreamProperties.StreamDistribution.SINGLE; import static com.facebook.presto.sql.planner.optimizations.StreamPropertyDerivations.derivePropertiesRecursively; @@ -719,6 +723,52 @@ private PlanWithProperties planTableWriteWithTableWriteMerge(TableWriterNode tab gatherExchangeWithProperties.getProperties()); } + private PlanWithProperties visitPartitionedWriter(PlanNode node, Optional optionalPartitioning, PlanNode source, StreamPreferredProperties parentPreferences) + { + if (getTaskWriterCount(session) == 1) { + return planAndEnforceChildren(node, singleStream(), defaultParallelism(session)); + } + + if (!optionalPartitioning.isPresent()) { + return planAndEnforceChildren(node, fixedParallelism(), fixedParallelism()); + } + + PartitioningScheme partitioningScheme = optionalPartitioning.get(); + + if (partitioningScheme.getPartitioning().getHandle().equals(FIXED_HASH_DISTRIBUTION)) { + // arbitrary hash function on predefined set of partition columns + StreamPreferredProperties preference = partitionedOn(partitioningScheme.getPartitioning().getVariableReferences()); + return planAndEnforceChildren(node, preference, preference); + } + + // connector provided hash function + verify(!(partitioningScheme.getPartitioning().getHandle().getConnectorHandle() instanceof SystemPartitioningHandle)); + // TODO #20578: Check if the following verification is correct. + verify(partitioningScheme.getPartitioning().getArguments().stream() + .noneMatch(argument -> argument instanceof ConstantExpression), + "Table writer partitioning has constant arguments"); + PlanWithProperties newSource = source.accept(this, parentPreferences); + PlanWithProperties exchange = deriveProperties( + partitionedExchange( + idAllocator.getNextId(), + LOCAL, + newSource.getNode(), + partitioningScheme), + newSource.getProperties()); + + return rebaseAndDeriveProperties(node, ImmutableList.of(exchange)); + } + + // + // Merge + // + + @Override + public PlanWithProperties visitMergeWriter(MergeWriterNode node, StreamPreferredProperties parentPreferences) + { + return visitPartitionedWriter(node, node.getPartitioningScheme(), node.getSource(), parentPreferences); + } + @Override public PlanWithProperties visitTableWriteMerge(TableWriterMergeNode node, StreamPreferredProperties context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java index b8b0089153b7b..0bbe7c1476f80 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java @@ -65,6 +65,8 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; @@ -426,6 +428,18 @@ public ActualProperties visitUpdate(UpdateNode node, List inpu return Iterables.getOnlyElement(inputProperties).translateVariable(symbol -> Optional.empty()); } + @Override + public ActualProperties visitMergeWriter(MergeWriterNode node, List inputProperties) + { + return visitPartitionedWriter(inputProperties); + } + + @Override + public ActualProperties visitMergeProcessor(MergeProcessorNode node, List inputProperties) + { + return Iterables.getOnlyElement(inputProperties).translateVariable(symbol -> Optional.empty()); + } + @Override public ActualProperties visitMetadataDelete(MetadataDeleteNode node, List inputProperties) { @@ -785,6 +799,11 @@ else if (!(value instanceof RowExpression)) { @Override public ActualProperties visitTableWriter(TableWriterNode node, List inputProperties) + { + return visitPartitionedWriter(inputProperties); + } + + private ActualProperties visitPartitionedWriter(List inputProperties) { ActualProperties properties = Iterables.getOnlyElement(inputProperties); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java index d9fd049555be3..e0027cdbba25e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java @@ -61,6 +61,8 @@ import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.SimplePlanRewriter; @@ -524,6 +526,49 @@ public PlanNode visitTableScan(TableScanNode node, RewriteContext> context) + { + Set expectedInputs = ImmutableSet.builder() + .addAll(node.getMergeProcessorProjectedVariables()) + .build(); + + PlanNode source = context.rewrite(node.getSource(), expectedInputs); + + return new MergeWriterNode( + node.getSourceLocation(), + node.getId(), + node.getStatsEquivalentPlanNode(), + source, + node.getTarget(), + node.getMergeProcessorProjectedVariables(), + node.getPartitioningScheme(), + node.getOutputVariables()); + } + + @Override + public PlanNode visitMergeProcessor(MergeProcessorNode node, RewriteContext> context) + { + Set expectedInputs = ImmutableSet.builder() + .add(node.getRowIdVariable()) + .add(node.getMergeRowVariable()) + .build(); + + PlanNode source = context.rewrite(node.getSource(), expectedInputs); + + return new MergeProcessorNode( + node.getSourceLocation(), + node.getId(), + node.getStatsEquivalentPlanNode(), + source, + node.getTarget(), + node.getRowIdVariable(), + node.getMergeRowVariable(), + node.getTargetColumnVariables(), + node.getTargetRedistributionColumnVariables(), + node.getOutputVariables()); + } + @Override public PlanNode visitFilter(FilterNode node, RewriteContext> context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java index 96288e6af729c..b1b44d241e6d6 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java @@ -105,6 +105,16 @@ public StreamPreferredProperties withFixedParallelism() return fixedParallelism(); } + public static StreamPreferredProperties partitionedOn(Collection partitionSymbols) + { + if (partitionSymbols.isEmpty()) { + return singleStream(); + } + + // Prefer partitioning on given partitioning symbols. Partition hash can be evaluated in any order. + return new StreamPreferredProperties(Optional.of(FIXED), false, Optional.of(ImmutableSet.copyOf(partitionSymbols)), false); + } + public static StreamPreferredProperties exactlyPartitionedOn(Collection partitionVariables) { if (partitionVariables.isEmpty()) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java index a56b3c773d6ed..3883a7f403b1a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java @@ -55,6 +55,8 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; @@ -482,6 +484,20 @@ public StreamProperties visitUpdate(UpdateNode node, List inpu return properties.withUnspecifiedPartitioning(); } + @Override + public StreamProperties visitMergeWriter(MergeWriterNode node, List inputProperties) + { + StreamProperties properties = Iterables.getOnlyElement(inputProperties); + return properties.withUnspecifiedPartitioning(); + } + + @Override + public StreamProperties visitMergeProcessor(MergeProcessorNode node, List inputProperties) + { + StreamProperties properties = Iterables.getOnlyElement(inputProperties); + return properties.withUnspecifiedPartitioning(); + } + @Override public StreamProperties visitMetadataDelete(MetadataDeleteNode node, List inputProperties) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/SymbolMapper.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/SymbolMapper.java index 9805efad17939..894359520b979 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/SymbolMapper.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/SymbolMapper.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.WarningCollector; import com.facebook.presto.spi.plan.AggregationNode; import com.facebook.presto.spi.plan.AggregationNode.Aggregation; +import com.facebook.presto.spi.plan.ExchangeEncoding; import com.facebook.presto.spi.plan.Ordering; import com.facebook.presto.spi.plan.OrderingScheme; import com.facebook.presto.spi.plan.PartitioningScheme; @@ -36,6 +37,8 @@ import com.facebook.presto.spi.relation.VariableReferenceExpression; import com.facebook.presto.sql.planner.Symbol; import com.facebook.presto.sql.planner.TypeProvider; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.tree.Expression; @@ -110,6 +113,13 @@ public VariableReferenceExpression map(VariableReferenceExpression variable) return new VariableReferenceExpression(variable.getSourceLocation(), canonical, types.get(new SymbolReference(getNodeLocation(variable.getSourceLocation()), canonical))); } + public List map(List variableReferenceExpressions) + { + return variableReferenceExpressions.stream() + .map(this::map) + .collect(toImmutableList()); + } + public Expression map(Expression value) { return ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter() @@ -274,6 +284,64 @@ public StatisticsWriterNode map(StatisticsWriterNode node, PlanNode source) node.getDescriptor().map(this::map)); } + public MergeWriterNode map(MergeWriterNode node, PlanNode source) + { + // Intentionally does not use mapAndDistinct on columns as that would remove columns + List newOutputs = map(node.getOutputVariables()); + + return new MergeWriterNode( + source.getSourceLocation(), + node.getId(), + source, + node.getTarget(), + map(node.getMergeProcessorProjectedVariables()), + node.getPartitioningScheme().map(partitioningScheme -> map(partitioningScheme, source.getOutputVariables())), + newOutputs); + } + + public MergeWriterNode map(MergeWriterNode node, PlanNode source, PlanNodeId newId) + { + // Intentionally does not use mapAndDistinct on columns as that would remove columns + List newOutputs = map(node.getOutputVariables()); + + return new MergeWriterNode( + source.getSourceLocation(), + newId, + source, + node.getTarget(), + map(node.getMergeProcessorProjectedVariables()), + node.getPartitioningScheme().map(partitioningScheme -> map(partitioningScheme, source.getOutputVariables())), + newOutputs); + } + + public MergeProcessorNode map(MergeProcessorNode node, PlanNode source) + { + List newOutputs = map(node.getOutputVariables()); + + return new MergeProcessorNode( + source.getSourceLocation(), + node.getId(), + source, + node.getTarget(), + map(node.getRowIdVariable()), + map(node.getMergeRowVariable()), + map(node.getTargetColumnVariables()), + map(node.getTargetRedistributionColumnVariables()), + newOutputs); + } + + public PartitioningScheme map(PartitioningScheme scheme, List sourceLayout) + { + return new PartitioningScheme( + translateVariable(scheme.getPartitioning(), this::map), + mapAndDistinctVariable(sourceLayout), + scheme.getHashColumn().map(this::map), + scheme.isReplicateNullsAndAny(), + scheme.isScaleWriters(), + ExchangeEncoding.COLUMNAR, + scheme.getBucketToPartition()); + } + public TableFinishNode map(TableFinishNode node, PlanNode source) { return new TableFinishNode( diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java index 91a92107f9f3c..c51ae28191e8e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java @@ -67,6 +67,8 @@ import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode; import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.OffsetNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; @@ -462,6 +464,50 @@ public PlanNode visitUpdate(UpdateNode node, RewriteContext context) return new UpdateNode(node.getSourceLocation(), node.getId(), node.getSource(), canonicalize(node.getRowId()), node.getColumnValueAndRowIdSymbols(), node.getOutputVariables()); } + @Override + public PlanNode visitMergeWriter(MergeWriterNode node, RewriteContext context) + { + // TODO #20578 Which implementation should we choose? A or B? + // A +// PlanNode source = context.rewrite(node.getSource()); +// SymbolMapper mapper = new SymbolMapper(mapping, types, warningCollector); +// return mapper.map(node, source); + + // B + return new MergeWriterNode( + node.getSourceLocation(), + node.getId(), + node.getStatsEquivalentPlanNode(), + context.rewrite(node.getSource()), + node.getTarget(), + node.getMergeProcessorProjectedVariables(), + node.getPartitioningScheme(), + node.getOutputVariables()); + } + + @Override + public PlanNode visitMergeProcessor(MergeProcessorNode node, RewriteContext context) + { + // TODO #20578 Which implementation should we choose? A or B? + // A +// PlanNode source = context.rewrite(node.getSource()); +// SymbolMapper mapper = new SymbolMapper(mapping, types, warningCollector); +// return mapper.map(node, source); + + // B + return new MergeProcessorNode( + node.getSourceLocation(), + node.getId(), + node.getStatsEquivalentPlanNode(), + context.rewrite(node.getSource()), + node.getTarget(), + node.getRowIdVariable(), + node.getMergeRowVariable(), + node.getTargetColumnVariables(), + node.getTargetRedistributionColumnVariables(), + node.getOutputVariables()); + } + @Override public PlanNode visitStatisticsWriterNode(StatisticsWriterNode node, RewriteContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java index b33dfc48938d7..507967484c7d6 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.sql.planner.plan; +import com.facebook.presto.spi.plan.MergeJoinNode; import com.facebook.presto.spi.plan.PlanVisitor; import com.facebook.presto.sql.planner.CanonicalJoinNode; import com.facebook.presto.sql.planner.CanonicalTableScanNode; @@ -37,6 +38,21 @@ public R visitExplainAnalyze(ExplainAnalyzeNode node, C context) return visitPlan(node, context); } + public R visitMergeJoin(MergeJoinNode node, C context) + { + return visitPlan(node, context); + } + + public R visitMergeWriter(MergeWriterNode node, C context) + { + return visitPlan(node, context); + } + + public R visitMergeProcessor(MergeProcessorNode node, C context) + { + return visitPlan(node, context); + } + public R visitOffset(OffsetNode node, C context) { return visitPlan(node, context); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeProcessorNode.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeProcessorNode.java new file mode 100644 index 0000000000000..77a9de3ca0773 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeProcessorNode.java @@ -0,0 +1,152 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner.plan; + +import com.facebook.presto.spi.SourceLocation; +import com.facebook.presto.spi.plan.PlanNode; +import com.facebook.presto.spi.plan.PlanNodeId; +import com.facebook.presto.spi.plan.TableWriterNode.MergeTarget; +import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; + +import java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * The node processes the result of the Searched CASE and RIGHT JOIN + * derived from a MERGE statement. + */ +public class MergeProcessorNode + extends InternalPlanNode +{ + private final PlanNode source; + private final MergeTarget target; + private final VariableReferenceExpression rowIdVariable; + private final VariableReferenceExpression mergeRowVariable; + private final List targetColumnVariables; + private final List targetRedistributionColumnVariables; + private final List outputs; + + @JsonCreator + public MergeProcessorNode( + Optional sourceLocation, + @JsonProperty("id") PlanNodeId id, + @JsonProperty("source") PlanNode source, + @JsonProperty("target") MergeTarget target, + @JsonProperty("rowIdVariable") VariableReferenceExpression rowIdVariable, + @JsonProperty("mergeRowVariable") VariableReferenceExpression mergeRowVariable, + @JsonProperty("targetColumnVariables") List targetColumnVariables, + @JsonProperty("targetRedistributionColumnVariables") List targetRedistributionColumnVariables, + @JsonProperty("outputs") List outputs) + { + this(sourceLocation, id, Optional.empty(), source, target, rowIdVariable, mergeRowVariable, targetColumnVariables, targetRedistributionColumnVariables, outputs); + } + + public MergeProcessorNode( + Optional sourceLocation, + PlanNodeId id, + Optional statsEquivalentPlanNode, + PlanNode source, + MergeTarget target, + VariableReferenceExpression rowIdVariable, + VariableReferenceExpression mergeRowVariable, + List targetColumnVariables, + List targetRedistributionColumnVariables, + List outputs) + { + super(sourceLocation, id, statsEquivalentPlanNode); + + this.source = requireNonNull(source, "source is null"); + this.target = requireNonNull(target, "target is null"); + this.mergeRowVariable = requireNonNull(mergeRowVariable, "mergeRowVariable is null"); + this.rowIdVariable = requireNonNull(rowIdVariable, "rowIdVariable is null"); + this.targetColumnVariables = requireNonNull(targetColumnVariables, "targetColumnVariables is null"); + this.targetRedistributionColumnVariables = requireNonNull(targetRedistributionColumnVariables, "targetRedistributionColumnVariables is null"); + this.outputs = ImmutableList.copyOf(requireNonNull(outputs, "outputs is null")); + } + + @JsonProperty + public PlanNode getSource() + { + return source; + } + + @JsonProperty + public MergeTarget getTarget() + { + return target; + } + + @JsonProperty + public VariableReferenceExpression getMergeRowVariable() + { + return mergeRowVariable; + } + + @JsonProperty + public VariableReferenceExpression getRowIdVariable() + { + return rowIdVariable; + } + + @JsonProperty + public List getTargetColumnVariables() + { + return targetColumnVariables; + } + + @JsonProperty + public List getTargetRedistributionColumnVariables() + { + return targetRedistributionColumnVariables; + } + + @JsonProperty("outputs") + @Override + public List getOutputVariables() + { + return outputs; + } + + @Override + public List getSources() + { + return ImmutableList.of(source); + } + + @Override + public R accept(InternalPlanVisitor visitor, C context) + { + return visitor.visitMergeProcessor(this, context); + } + + @Override + public PlanNode replaceChildren(List newChildren) + { + return new MergeProcessorNode(getSourceLocation(), getId(), Iterables.getOnlyElement(newChildren), + target, rowIdVariable, mergeRowVariable, targetColumnVariables, targetRedistributionColumnVariables, outputs); + } + + @Override + public PlanNode assignStatsEquivalentPlanNode(Optional statsEquivalentPlanNode) + { + return new MergeProcessorNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, target, + rowIdVariable, mergeRowVariable, targetColumnVariables, targetRedistributionColumnVariables, outputs); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeWriterNode.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeWriterNode.java new file mode 100644 index 0000000000000..28b397aa9ad4a --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/MergeWriterNode.java @@ -0,0 +1,131 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner.plan; + +import com.facebook.presto.spi.SourceLocation; +import com.facebook.presto.spi.plan.PartitioningScheme; +import com.facebook.presto.spi.plan.PlanNode; +import com.facebook.presto.spi.plan.PlanNodeId; +import com.facebook.presto.spi.plan.TableWriterNode.MergeTarget; +import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.errorprone.annotations.Immutable; + +import java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +@Immutable +public class MergeWriterNode + extends InternalPlanNode +{ + private final PlanNode source; + private final MergeTarget target; + private final List mergeProcessorProjectedVariables; + private final Optional partitioningScheme; + private final List outputs; + + @JsonCreator + public MergeWriterNode( + Optional sourceLocation, + @JsonProperty("id") PlanNodeId id, + @JsonProperty("source") PlanNode source, + @JsonProperty("target") MergeTarget target, + @JsonProperty("mergeProcessorProjectedVariables") List mergeProcessorProjectedVariables, + @JsonProperty("partitioningScheme") Optional partitioningScheme, + @JsonProperty("outputs") List outputs) + { + this(sourceLocation, id, Optional.empty(), source, target, mergeProcessorProjectedVariables, partitioningScheme, outputs); + } + + public MergeWriterNode( + Optional sourceLocation, + PlanNodeId id, + Optional statsEquivalentPlanNode, + PlanNode source, + MergeTarget target, + List mergeProcessorProjectedVariables, + Optional partitioningScheme, + List outputs) + { + super(sourceLocation, id, statsEquivalentPlanNode); + + this.source = requireNonNull(source, "source is null"); + this.target = requireNonNull(target, "target is null"); + this.mergeProcessorProjectedVariables = requireNonNull(mergeProcessorProjectedVariables, "mergeProcessorProjectedVariables is null"); + this.partitioningScheme = requireNonNull(partitioningScheme, "partitioningScheme is null"); + this.outputs = ImmutableList.copyOf(requireNonNull(outputs, "outputs is null")); + } + + @JsonProperty + public PlanNode getSource() + { + return source; + } + + @JsonProperty + public MergeTarget getTarget() + { + return target; + } + + @JsonProperty + public List getMergeProcessorProjectedVariables() + { + return mergeProcessorProjectedVariables; + } + + @JsonProperty + public Optional getPartitioningScheme() + { + return partitioningScheme; + } + + @Override + public List getSources() + { + return ImmutableList.of(source); + } + + @Override + @JsonProperty("outputs") + public List getOutputVariables() + { + return outputs; + } + + @Override + public R accept(InternalPlanVisitor visitor, C context) + { + return visitor.visitMergeWriter(this, context); + } + + @Override + public PlanNode assignStatsEquivalentPlanNode(Optional statsEquivalentPlanNode) + { + return new MergeWriterNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, target, + mergeProcessorProjectedVariables, partitioningScheme, outputs); + } + + @Override + public PlanNode replaceChildren(List newChildren) + { + return new MergeWriterNode(getSourceLocation(), getId(), Iterables.getOnlyElement(newChildren), + target, mergeProcessorProjectedVariables, partitioningScheme, outputs); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java index f1a00c4b1a128..25c3e345e3793 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java @@ -204,6 +204,11 @@ public static Pattern tableWriterMergeNode() return typeOf(TableWriterMergeNode.class); } + public static Pattern mergeWriter() + { + return typeOf(MergeWriterNode.class); + } + public static Pattern topN() { return typeOf(TopNNode.class); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java index 2b7059d12e02a..4fc8999032e7f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java @@ -92,6 +92,8 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.OffsetNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; @@ -1273,6 +1275,26 @@ public Void visitMetadataDelete(MetadataDeleteNode node, Void context) return processChildren(node, context); } + @Override + public Void visitMergeWriter(MergeWriterNode node, Void context) + { + addNode(node, "MergeWriter", format("table: %s", node.getTarget().toString())); + return processChildren(node, context); + } + + @Override + public Void visitMergeProcessor(MergeProcessorNode node, Void context) + { + NodeRepresentation nodeOutput = addNode(node, "MergeProcessor"); + nodeOutput.appendDetails("target: %s", node.getTarget()); + nodeOutput.appendDetails("merge row column: %s", node.getMergeRowVariable()); + nodeOutput.appendDetails("row id column: %s", node.getRowIdVariable()); + nodeOutput.appendDetails("redistribution columns: %s", node.getTargetRedistributionColumnVariables()); + nodeOutput.appendDetails("data columns: %s", node.getTargetColumnVariables()); + + return processChildren(node, context); + } + @Override public Void visitEnforceSingleRow(EnforceSingleRowNode node, Void context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java index 7bb6f516b9171..c2529a0c8280e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java @@ -63,6 +63,8 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.OffsetNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; @@ -613,6 +615,30 @@ public Void visitTableWriteMerge(TableWriterMergeNode node, Set boundSymbols) + { + PlanNode source = node.getSource(); + source.accept(this, boundSymbols); // visit child + return null; + } + + @Override + public Void visitMergeProcessor(MergeProcessorNode node, Set boundSymbols) + { + PlanNode source = node.getSource(); + source.accept(this, boundSymbols); // visit child + + checkArgument(source.getOutputVariables().contains(node.getRowIdVariable()), + "Invalid node. rowId symbol (%s) is not in source plan output (%s)", + node.getRowIdVariable(), node.getSource().getOutputVariables()); + checkArgument(source.getOutputVariables().contains(node.getMergeRowVariable()), + "Invalid node. Merge row symbol (%s) is not in source plan output (%s)", + node.getMergeRowVariable(), node.getSource().getOutputVariables()); + + return null; + } + @Override public Void visitDelete(DeleteNode node, Set boundVariables) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/util/GraphvizPrinter.java b/presto-main-base/src/main/java/com/facebook/presto/util/GraphvizPrinter.java index 6c210e9e0848c..b74924840394f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/util/GraphvizPrinter.java +++ b/presto-main-base/src/main/java/com/facebook/presto/util/GraphvizPrinter.java @@ -61,6 +61,8 @@ import com.facebook.presto.sql.planner.plan.GroupIdNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.LateralJoinNode; +import com.facebook.presto.sql.planner.plan.MergeProcessorNode; +import com.facebook.presto.sql.planner.plan.MergeWriterNode; import com.facebook.presto.sql.planner.plan.RemoteSourceNode; import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; @@ -131,6 +133,7 @@ private enum NodeType ANALYZE_FINISH, EXPLAIN_ANALYZE, UPDATE, + MERGE } private static final Map NODE_COLORS = immutableEnumMap(ImmutableMap.builder() @@ -162,6 +165,7 @@ private enum NodeType .put(NodeType.ANALYZE_FINISH, "plum") .put(NodeType.EXPLAIN_ANALYZE, "cadetblue1") .put(NodeType.UPDATE, "blue") + .put(NodeType.MERGE, "lightblue") .build()); static { @@ -321,6 +325,20 @@ public Void visitUpdate(UpdateNode node, Void context) return node.getSource().accept(this, context); } + @Override + public Void visitMergeWriter(MergeWriterNode node, Void context) + { + printNode(node, format("MergeWriterNode[%s]", Joiner.on(", ").join(node.getOutputVariables())), NODE_COLORS.get(NodeType.MERGE)); + return node.getSource().accept(this, context); + } + + @Override + public Void visitMergeProcessor(MergeProcessorNode node, Void context) + { + printNode(node, format("MergeProcessorNode[%s]", Joiner.on(", ").join(node.getOutputVariables())), NODE_COLORS.get(NodeType.MERGE)); + return node.getSource().accept(this, context); + } + @Override public Void visitStatisticsWriterNode(StatisticsWriterNode node, Void context) { diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/exchange/TestLocalExchange.java b/presto-main-base/src/test/java/com/facebook/presto/operator/exchange/TestLocalExchange.java index 465a374eca960..06fb18c1502d1 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/exchange/TestLocalExchange.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/exchange/TestLocalExchange.java @@ -460,9 +460,9 @@ public void testCreatePartitionFunction() new ConnectorId("prism"), new ConnectorNodePartitioningProvider() { @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { - return createBucketNodeMap(Stream.generate(() -> sortedNodes).flatMap(List::stream).limit(10).collect(toImmutableList()), SOFT_AFFINITY); + return Optional.of(createBucketNodeMap(Stream.generate(() -> sortedNodes).flatMap(List::stream).limit(10).collect(toImmutableList()), SOFT_AFFINITY)); } @Override diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp index efb585849ef31..a69a62d38af7f 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp @@ -2288,6 +2288,10 @@ void to_json(json& j, const std::shared_ptr& p) { j = *std::static_pointer_cast(p); return; } + if (type == "MergeHandle") { + j = *std::static_pointer_cast(p); + return; + } throw TypeError(type + " no abstract type ExecutionWriterTarget "); } @@ -2332,6 +2336,12 @@ void from_json(const json& j, std::shared_ptr& p) { p = std::static_pointer_cast(k); return; } + if (type == "MergeHandle") { + std::shared_ptr k = std::make_shared(); + j.get_to(*k); + p = std::static_pointer_cast(k); + return; + } throw TypeError(type + " no abstract type ExecutionWriterTarget "); } @@ -6952,6 +6962,71 @@ void from_json(const json& j, MemoryInfo& p) { } } // namespace facebook::presto::protocol namespace facebook::presto::protocol { +void to_json(json& j, const std::shared_ptr& p) { + if (p == nullptr) { + return; + } + String type = p->_type; + + throw TypeError(type + " no abstract type ConnectorMergeTableHandle "); +} + +void from_json(const json& j, std::shared_ptr& p) { + String type; + try { + type = p->getSubclassKey(j); + } catch (json::parse_error& e) { + throw ParseError( + std::string(e.what()) + + " ConnectorMergeTableHandle ConnectorMergeTableHandle"); + } + + throw TypeError(type + " no abstract type ConnectorMergeTableHandle "); +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +MergeHandle::MergeHandle() noexcept { + _type = "MergeHandle"; +} + +void to_json(json& j, const MergeHandle& p) { + j = json::object(); + j["@type"] = "MergeHandle"; + to_json_key( + j, + "tableHandle", + p.tableHandle, + "MergeHandle", + "TableHandle", + "tableHandle"); + to_json_key( + j, + "connectorMergeTableHandle", + p.connectorMergeTableHandle, + "MergeHandle", + "ConnectorMergeTableHandle", + "connectorMergeTableHandle"); +} + +void from_json(const json& j, MergeHandle& p) { + p._type = j["@type"]; + from_json_key( + j, + "tableHandle", + p.tableHandle, + "MergeHandle", + "TableHandle", + "tableHandle"); + from_json_key( + j, + "connectorMergeTableHandle", + p.connectorMergeTableHandle, + "MergeHandle", + "ConnectorMergeTableHandle", + "connectorMergeTableHandle"); +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { MergeJoinNode::MergeJoinNode() noexcept { _type = ".MergeJoinNode"; } @@ -7044,6 +7119,151 @@ void from_json(const json& j, MergeJoinNode& p) { } } // namespace facebook::presto::protocol namespace facebook::presto::protocol { +// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM() + +// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays +static const std::pair RowChangeParadigm_enum_table[] = + { // NOLINT: cert-err58-cpp + {RowChangeParadigm::CHANGE_ONLY_UPDATED_COLUMNS, + "CHANGE_ONLY_UPDATED_COLUMNS"}, + {RowChangeParadigm::DELETE_ROW_AND_INSERT_ROW, + "DELETE_ROW_AND_INSERT_ROW"}}; +void to_json(json& j, const RowChangeParadigm& e) { + static_assert( + std::is_enum::value, + "RowChangeParadigm must be an enum!"); + const auto* it = std::find_if( + std::begin(RowChangeParadigm_enum_table), + std::end(RowChangeParadigm_enum_table), + [e](const std::pair& ej_pair) -> bool { + return ej_pair.first == e; + }); + j = ((it != std::end(RowChangeParadigm_enum_table)) + ? it + : std::begin(RowChangeParadigm_enum_table)) + ->second; +} +void from_json(const json& j, RowChangeParadigm& e) { + static_assert( + std::is_enum::value, + "RowChangeParadigm must be an enum!"); + const auto* it = std::find_if( + std::begin(RowChangeParadigm_enum_table), + std::end(RowChangeParadigm_enum_table), + [&j](const std::pair& ej_pair) -> bool { + return ej_pair.second == j; + }); + e = ((it != std::end(RowChangeParadigm_enum_table)) + ? it + : std::begin(RowChangeParadigm_enum_table)) + ->first; +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { + +void to_json(json& j, const MergeParadigmAndTypes& p) { + j = json::object(); + to_json_key( + j, + "paradigm", + p.paradigm, + "MergeParadigmAndTypes", + "RowChangeParadigm", + "paradigm"); + to_json_key( + j, + "columnTypes", + p.columnTypes, + "MergeParadigmAndTypes", + "List", + "columnTypes"); + to_json_key( + j, + "rowIdType", + p.rowIdType, + "MergeParadigmAndTypes", + "Type", + "rowIdType"); +} + +void from_json(const json& j, MergeParadigmAndTypes& p) { + from_json_key( + j, + "paradigm", + p.paradigm, + "MergeParadigmAndTypes", + "RowChangeParadigm", + "paradigm"); + from_json_key( + j, + "columnTypes", + p.columnTypes, + "MergeParadigmAndTypes", + "List", + "columnTypes"); + from_json_key( + j, + "rowIdType", + p.rowIdType, + "MergeParadigmAndTypes", + "Type", + "rowIdType"); +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { + +void to_json(json& j, const MergeTarget& p) { + j = json::object(); + to_json_key(j, "handle", p.handle, "MergeTarget", "TableHandle", "handle"); + to_json_key( + j, + "mergeHandle", + p.mergeHandle, + "MergeTarget", + "MergeHandle", + "mergeHandle"); + to_json_key( + j, + "schemaTableName", + p.schemaTableName, + "MergeTarget", + "SchemaTableName", + "schemaTableName"); + to_json_key( + j, + "mergeParadigmAndTypes", + p.mergeParadigmAndTypes, + "MergeTarget", + "MergeParadigmAndTypes", + "mergeParadigmAndTypes"); +} + +void from_json(const json& j, MergeTarget& p) { + from_json_key(j, "handle", p.handle, "MergeTarget", "TableHandle", "handle"); + from_json_key( + j, + "mergeHandle", + p.mergeHandle, + "MergeTarget", + "MergeHandle", + "mergeHandle"); + from_json_key( + j, + "schemaTableName", + p.schemaTableName, + "MergeTarget", + "SchemaTableName", + "schemaTableName"); + from_json_key( + j, + "mergeParadigmAndTypes", + p.mergeParadigmAndTypes, + "MergeTarget", + "MergeParadigmAndTypes", + "mergeParadigmAndTypes"); +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { void to_json(json& j, const NodeLoadMetrics& p) { j = json::object(); diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h index 2b1e4eb66c14e..029b2876d188b 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h @@ -319,6 +319,11 @@ struct ColumnHandle : public JsonEncodedSubclass { void to_json(json& j, const std::shared_ptr& p); void from_json(const json& j, std::shared_ptr& p); } // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +struct ConnectorMergeTableHandle : public JsonEncodedSubclass {}; +void to_json(json& j, const std::shared_ptr& p); +void from_json(const json& j, std::shared_ptr& p); +} // namespace facebook::presto::protocol namespace facebook::presto::protocol { struct SourceLocation { @@ -1745,6 +1750,16 @@ void to_json(json& j, const MemoryInfo& p); void from_json(const json& j, MemoryInfo& p); } // namespace facebook::presto::protocol namespace facebook::presto::protocol { +struct MergeHandle : public ExecutionWriterTarget { + TableHandle tableHandle = {}; + std::shared_ptr connectorMergeTableHandle = {}; + + MergeHandle() noexcept; +}; +void to_json(json& j, const MergeHandle& p); +void from_json(const json& j, MergeHandle& p); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { struct MergeJoinNode : public PlanNode { MergeJoinNode() noexcept; PlanNodeId id = {}; @@ -1761,6 +1776,33 @@ void to_json(json& j, const MergeJoinNode& p); void from_json(const json& j, MergeJoinNode& p); } // namespace facebook::presto::protocol namespace facebook::presto::protocol { +enum class RowChangeParadigm { + CHANGE_ONLY_UPDATED_COLUMNS, + DELETE_ROW_AND_INSERT_ROW +}; +extern void to_json(json& j, const RowChangeParadigm& e); +extern void from_json(const json& j, RowChangeParadigm& e); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +struct MergeParadigmAndTypes { + RowChangeParadigm paradigm = {}; + List columnTypes = {}; + Type rowIdType = {}; +}; +void to_json(json& j, const MergeParadigmAndTypes& p); +void from_json(const json& j, MergeParadigmAndTypes& p); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +struct MergeTarget { + TableHandle handle = {}; + std::shared_ptr mergeHandle = {}; + SchemaTableName schemaTableName = {}; + MergeParadigmAndTypes mergeParadigmAndTypes = {}; +}; +void to_json(json& j, const MergeTarget& p); +void from_json(const json& j, MergeTarget& p); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { struct NodeLoadMetrics { double cpuUsedPercent = {}; double memoryUsedInBytes = {}; diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml index 9d930555c62ba..602c3d390b3de 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml @@ -85,6 +85,10 @@ AbstractClasses: super: JsonEncodedSubclass subclasses: + ConnectorMergeTableHandle: + super: JsonEncodedSubclass + subclasses: + ConnectorTransactionHandle: super: JsonEncodedSubclass subclasses: @@ -129,6 +133,7 @@ AbstractClasses: - { name: InsertHandle, key: InsertHandle } - { name: DeleteHandle, key: DeleteHandle } - { name: UpdateHandle, key: UpdateHandle } + - { name: MergeHandle, key: MergeHandle } InputDistribution: super: JsonEncodedSubclass @@ -220,6 +225,7 @@ JavaClasses: - presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSplit.java - presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableHandle.java - presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorTransactionHandle.java + - presto-spi/src/main/java/com/facebook/presto/spi/connector/RowChangeParadigm.java - presto-spi/src/main/java/com/facebook/presto/spi/ConnectorIndexHandle.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/DistinctLimitNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/MarkDistinctNode.java @@ -343,6 +349,7 @@ JavaClasses: - presto-function-namespace-managers-common/src/main/java/com/facebook/presto/functionNamespace/JsonBasedUdfFunctionMetadata.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/DeleteNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/BaseInputDistribution.java + - presto-spi/src/main/java/com/facebook/presto/spi/MergeHandle.java - presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java - presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java - presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java index 38142dfcf3d6c..730e71205ce0e 100644 --- a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java +++ b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java @@ -24,19 +24,20 @@ import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import java.util.List; +import java.util.Optional; import java.util.function.ToIntFunction; public class PinotNodePartitioningProvider implements ConnectorNodePartitioningProvider { @Override - public ConnectorBucketNodeMap getBucketNodeMap( + public Optional getBucketNodeMap( ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { - return ConnectorBucketNodeMap.createBucketNodeMap(1); + return Optional.of(ConnectorBucketNodeMap.createBucketNodeMap(1)); } @Override diff --git a/presto-plugin-toolkit/src/main/java/com/facebook/presto/plugin/base/util/Closables.java b/presto-plugin-toolkit/src/main/java/com/facebook/presto/plugin/base/util/Closables.java new file mode 100644 index 0000000000000..d6293631808e9 --- /dev/null +++ b/presto-plugin-toolkit/src/main/java/com/facebook/presto/plugin/base/util/Closables.java @@ -0,0 +1,41 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.plugin.base.util; + +import static java.util.Objects.requireNonNull; + +public final class Closables +{ + private Closables() {} + + public static T closeAllSuppress(T rootCause, AutoCloseable... closeables) + { + requireNonNull(rootCause, "rootCause is null"); + requireNonNull(closeables, "closeables is null"); + for (AutoCloseable closeable : closeables) { + try { + if (closeable != null) { + closeable.close(); + } + } + catch (Throwable e) { + // Self-suppression not permitted + if (rootCause != e) { + rootCause.addSuppressed(e); + } + } + } + return rootCause; + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorHandleResolver.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorHandleResolver.java index cded59a56df0b..fe7ef06c58f46 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorHandleResolver.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorHandleResolver.java @@ -46,6 +46,11 @@ default Class getDeleteTableHandleClass() throw new UnsupportedOperationException(); } + default Class getMergeTableHandleClass() + { + throw new UnsupportedOperationException(); + } + default Class getPartitioningHandleClass() { throw new UnsupportedOperationException(); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeSink.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeSink.java new file mode 100644 index 0000000000000..aa78542002c49 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeSink.java @@ -0,0 +1,61 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi; + +import com.facebook.presto.common.Page; +import io.airlift.slice.Slice; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +public interface ConnectorMergeSink +{ + /** + * Represents an inserted row. + */ + int INSERT_OPERATION_NUMBER = 1; + + /** + * Represents a deleted row. + */ + int DELETE_OPERATION_NUMBER = 2; + + /** + * Represents an updated row when using {@link com.facebook.presto.spi.connector.RowChangeParadigm#CHANGE_ONLY_UPDATED_COLUMNS}. + */ + int UPDATE_OPERATION_NUMBER = 3; + + /** + * Store the page resulting from a merge. The page consists of {@code n} channels, numbered {@code 0..n-1}: + *

    + *
  • Blocks {@code 0..n-3} in page are the data columns
    • + *
  • Block {@code n-2} is the tinyint operation: + *
      + *
    • {@link #INSERT_OPERATION_NUMBER}
      • + *
    • {@link #DELETE_OPERATION_NUMBER}
      • + *
    • {@link #UPDATE_OPERATION_NUMBER}
      • + *
    + *
  • Block {@code n-1} is a connector-specific rowId column, whose handle was previously returned by + * {@link com.facebook.presto.spi.connector.ConnectorMetadata#getMergeRowIdColumnHandle(ConnectorSession, ConnectorTableHandle) getMergeRowIdColumnHandle()} + *
    • + *
+ * + * @param page The page to store. + */ + void storeMergedRows(Page page); + + CompletableFuture> finish(); + + default void abort() {} +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeTableHandle.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeTableHandle.java new file mode 100644 index 0000000000000..117a81841baa3 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorMergeTableHandle.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi; + +import com.facebook.presto.spi.connector.ConnectorMetadata; + +public interface ConnectorMergeTableHandle +{ + /** + * This method is required because the {@link ConnectorTableHandle} returned by + * {@link ConnectorMetadata#beginMerge} is in general different from the + * one passed to that method, but the updated handle must be made + * available to {@link ConnectorMetadata#finishMerge} + * + * @return the {@link ConnectorTableHandle} returned by {@link ConnectorMetadata#beginMerge} + */ + ConnectorTableHandle getTableHandle(); +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorNewTableLayout.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorNewTableLayout.java index 46fcf1eb665f8..652778d09ff35 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorNewTableLayout.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorNewTableLayout.java @@ -17,13 +17,14 @@ import java.util.List; import java.util.Objects; +import java.util.Optional; import static com.facebook.presto.spi.PartitionedTableWritePolicy.SINGLE_WRITER_PER_PARTITION_REQUIRED; import static java.util.Objects.requireNonNull; public class ConnectorNewTableLayout { - private final ConnectorPartitioningHandle partitioning; + private final Optional partitioning; private final List partitionColumns; private final PartitionedTableWritePolicy writerPolicy; @@ -34,12 +35,19 @@ public ConnectorNewTableLayout(ConnectorPartitioningHandle partitioning, List partitionColumns, PartitionedTableWritePolicy writerPolicy) { - this.partitioning = requireNonNull(partitioning, "partitioning is null"); + this.partitioning = Optional.of(requireNonNull(partitioning, "partitioning is null")); this.partitionColumns = requireNonNull(partitionColumns, "partitionColumns is null"); this.writerPolicy = requireNonNull(writerPolicy, "writerPolicy is null"); } - public ConnectorPartitioningHandle getPartitioning() + public ConnectorNewTableLayout(List partitionColumns) + { + this.partitioning = Optional.empty(); + this.partitionColumns = requireNonNull(partitionColumns, "partitionColumns is null"); + this.writerPolicy = SINGLE_WRITER_PER_PARTITION_REQUIRED; + } + + public Optional getPartitioning() { return partitioning; } diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/MergeHandle.java b/presto-spi/src/main/java/com/facebook/presto/spi/MergeHandle.java new file mode 100644 index 0000000000000..b9bd39a73eb98 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/MergeHandle.java @@ -0,0 +1,74 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +import static java.util.Objects.requireNonNull; + +public final class MergeHandle +{ + private final TableHandle tableHandle; + private final ConnectorMergeTableHandle connectorMergeTableHandle; + + @JsonCreator + public MergeHandle( + @JsonProperty("tableHandle") TableHandle tableHandle, + @JsonProperty("connectorMergeTableHandle") ConnectorMergeTableHandle connectorMergeTableHandle) + { + this.tableHandle = requireNonNull(tableHandle, "tableHandle is null"); + this.connectorMergeTableHandle = requireNonNull(connectorMergeTableHandle, "connectorMergeTableHandle is null"); + } + + @JsonProperty + public TableHandle getTableHandle() + { + return tableHandle; + } + + @JsonProperty + public ConnectorMergeTableHandle getConnectorMergeTableHandle() + { + return connectorMergeTableHandle; + } + + @Override + public int hashCode() + { + return Objects.hash(tableHandle, connectorMergeTableHandle); + } + + @Override + public boolean equals(Object obj) + { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + MergeHandle o = (MergeHandle) obj; + return Objects.equals(this.tableHandle, o.tableHandle) && + Objects.equals(this.connectorMergeTableHandle, o.connectorMergeTableHandle); + } + + @Override + public String toString() + { + return tableHandle + ":" + connectorMergeTableHandle; + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NewTableLayout.java b/presto-spi/src/main/java/com/facebook/presto/spi/NewTableLayout.java index 202b25c9a4923..93761adcac26d 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/NewTableLayout.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/NewTableLayout.java @@ -54,9 +54,10 @@ public ConnectorNewTableLayout getLayout() return layout; } - public PartitioningHandle getPartitioning() + public Optional getPartitioning() { - return new PartitioningHandle(Optional.of(connectorId), Optional.of(transactionHandle), layout.getPartitioning()); + return layout.getPartitioning() + .map(partitioning -> new PartitioningHandle(Optional.of(connectorId), Optional.of(transactionHandle), partitioning)); } public List getPartitionColumns() diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/TableHandle.java b/presto-spi/src/main/java/com/facebook/presto/spi/TableHandle.java index 3dd4c8abfad8e..a38055fbaacf2 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/TableHandle.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/TableHandle.java @@ -72,6 +72,16 @@ public TableHandle( this.dynamicFilter = requireNonNull(dynamicFilter, "dynamicFilter is null"); } + public TableHandle cloneWithConnectorHandle(ConnectorTableHandle connectorHandle) + { + return new TableHandle( + connectorId, + connectorHandle, + transaction, + layout, + dynamicFilter); + } + @JsonProperty @ThriftField(1) public ConnectorId getConnectorId() diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorMetadata.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorMetadata.java index e8a888bbc8ab9..4c9d5428d817a 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorMetadata.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorMetadata.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorNewTableLayout; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorResolvedIndex; @@ -71,6 +72,8 @@ public interface ConnectorMetadata { + String MODIFYING_ROWS_MESSAGE = "This connector does not support modifying table rows"; + /** * Checks if a schema exists. The connector may have schemas that exist * but are not enumerable via {@link #listSchemaNames}. @@ -561,6 +564,16 @@ default Optional getUpdateRowIdColumn(ConnectorSession session, Co return Optional.ofNullable(getUpdateRowIdColumnHandle(session, tableHandle, updatedColumns)); } + /** + * Get the column handle that will generate row IDs for the merge operation. + * These IDs will be passed to the {@link com.facebook.presto.spi.ConnectorMergeSink#storeMergedRows} + * method of the {@link com.facebook.presto.spi.ConnectorMergeSink} that created them. + */ + default ColumnHandle getMergeRowIdColumnHandle(ConnectorSession session, ConnectorTableHandle tableHandle) + { + throw new PrestoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + /** * Begin delete query */ @@ -600,6 +613,46 @@ default void finishUpdate(ConnectorSession session, ConnectorTableHandle tableHa throw new PrestoException(NOT_SUPPORTED, "This connector does not support update"); } + /** + * Return the row change paradigm supported by the connector on the table. + */ + default RowChangeParadigm getRowChangeParadigm(ConnectorSession session, ConnectorTableHandle tableHandle) + { + throw new PrestoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + + /** + * Get the physical layout for updated rows of a MERGE operation. + * Inserted rows are handled by {@link #getInsertLayout}. + * This layout always uses the {@link #getMergeRowIdColumnHandle merge row ID column}. + */ + default Optional getMergeUpdateLayout(ConnectorSession session, ConnectorTableHandle tableHandle) + { + return Optional.empty(); + } + + /** + * Do whatever is necessary to start an MERGE query, returning the {@link ConnectorMergeTableHandle} + * instance that will be passed to the PageSink, and to the {@link #finishMerge} method. + */ + default ConnectorMergeTableHandle beginMerge(ConnectorSession session, ConnectorTableHandle tableHandle) + { + throw new PrestoException(NOT_SUPPORTED, MODIFYING_ROWS_MESSAGE); + } + + /** + * Finish a merge query + * + * @param session The session + * @param mergeTableHandle A ConnectorMergeTableHandle for the table that is the target of the merge + * @param fragments All fragments returned by the merge plan + * @param computedStatistics Statistics for the table, meaningful only to the connector that produced them. + */ + default void finishMerge(ConnectorSession session, ConnectorMergeTableHandle mergeTableHandle, Collection fragments, Collection computedStatistics) + { + throw new PrestoException(GENERIC_INTERNAL_ERROR, "ConnectorMetadata beginMerge() is implemented without finishMerge()"); + } + /** * Create the specified view. The data for the view is opaque to the connector. */ diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java index bf35da6986437..c202d75e9be3e 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.Node; import java.util.List; +import java.util.Optional; import java.util.function.ToIntFunction; import static com.facebook.presto.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED; @@ -44,7 +45,7 @@ default List listPartitionHandles( return singletonList(NOT_PARTITIONED); } - ConnectorBucketNodeMap getBucketNodeMap( + Optional getBucketNodeMap( ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorPageSinkProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorPageSinkProvider.java index 5b8665d0f14fc..b3c1aa76ec8bf 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorPageSinkProvider.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorPageSinkProvider.java @@ -15,10 +15,15 @@ import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeSink; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorPageSink; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.PageSinkContext; +import com.facebook.presto.spi.PrestoException; + +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; public interface ConnectorPageSinkProvider { @@ -30,4 +35,9 @@ default ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionH { throw new UnsupportedOperationException("ConnectorPageSinkProvider does not support connectorDeleteTableHandle"); } + + default ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorMergeTableHandle mergeHandle) + { + throw new PrestoException(NOT_SUPPORTED, "This connector does not support SQL MERGE operations"); + } } diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/MergePage.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/MergePage.java new file mode 100644 index 0000000000000..e6c01c3f66319 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/MergePage.java @@ -0,0 +1,128 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.presto.spi.connector; + +import com.facebook.presto.common.Page; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.spi.ConnectorMergeSink; + +import java.util.Optional; +import java.util.stream.IntStream; + +import static com.facebook.presto.common.type.TinyintType.TINYINT; +import static com.facebook.presto.spi.ConnectorMergeSink.DELETE_OPERATION_NUMBER; +import static com.facebook.presto.spi.ConnectorMergeSink.INSERT_OPERATION_NUMBER; +import static java.lang.Math.toIntExact; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +/** + * Separate deletions and insertions pages from a merge using + * {@link RowChangeParadigm#DELETE_ROW_AND_INSERT_ROW}. + */ +public final class MergePage +{ + private final Optional deletionsPage; + private final Optional insertionsPage; + + private MergePage(Optional deletionsPage, Optional insertionsPage) + { + this.deletionsPage = requireNonNull(deletionsPage); + this.insertionsPage = requireNonNull(insertionsPage); + } + + /** + * @return delete page with data columns followed by row ID column + */ + public Optional getDeletionsPage() + { + return deletionsPage; + } + + /** + * @return insert page with data columns + */ + public Optional getInsertionsPage() + { + return insertionsPage; + } + + /** + * @param inputPage It has N + 2 channels/blocks, where N is the number of columns in the source table.
+ * 1: Source table column 1.
+ * 2: Source table column 2.
+ * N: Source table column N.
+ * N + 1: Operation: INSERT(1), DELETE(2), UPDATE(3). More info: {@link ConnectorMergeSink}
+ * N + 2: Merge Row ID (_file:varchar, _pos:bigint, file_record_count:bigint, partition_spec_id:integer, partition_data:varchar).
+ * @param dataColumnCount Number of columns of the MERGE INTO target table. + */ + public static MergePage createDeleteAndInsertPages(Page inputPage, int dataColumnCount) + { + // see page description in ConnectorMergeSink + int inputChannelCount = inputPage.getChannelCount(); + if (inputChannelCount != dataColumnCount + 2) { + throw new IllegalArgumentException(format("inputPage channelCount (%s) == dataColumns size (%s) + 2", inputChannelCount, dataColumnCount)); + } + + int positionCount = inputPage.getPositionCount(); // number of rows inserted, deleted or updated in this page. The updated rows count double. + if (positionCount <= 0) { + throw new IllegalArgumentException("positionCount should be > 0, but is " + positionCount); + } + + Block operationBlock = inputPage.getBlock(inputChannelCount - 2); + + int[] deletePositions = new int[positionCount]; + int[] insertPositions = new int[positionCount]; + int deletePositionCount = 0; + int insertPositionCount = 0; + + for (int position = 0; position < positionCount; position++) { + int operation = toIntExact(TINYINT.getLong(operationBlock, position)); + switch (operation) { + case DELETE_OPERATION_NUMBER: + deletePositions[deletePositionCount] = position; + deletePositionCount++; + break; + case INSERT_OPERATION_NUMBER: + insertPositions[insertPositionCount] = position; + insertPositionCount++; + break; + default: + throw new IllegalArgumentException("Invalid merge operation: " + operation); + } + } + + Optional deletePage = Optional.empty(); + if (deletePositionCount > 0) { + int[] columns = new int[dataColumnCount + 1]; + for (int i = 0; i < dataColumnCount; i++) { + columns[i] = i; + } + columns[dataColumnCount] = dataColumnCount + 1; // Merge Row ID channel + deletePage = Optional.of(inputPage + .getColumns(columns) + .getPositions(deletePositions, 0, deletePositionCount)); + } + + Optional insertPage = Optional.empty(); + if (insertPositionCount > 0) { + insertPage = Optional.of(inputPage + .getColumns(IntStream.range(0, dataColumnCount).toArray()) + .getPositions(insertPositions, 0, insertPositionCount)); + } + + return new MergePage(deletePage, insertPage); + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/RowChangeParadigm.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/RowChangeParadigm.java new file mode 100644 index 0000000000000..4ab6ef9a1bb0c --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/RowChangeParadigm.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi.connector; + +/** + * Different connectors have different ways of representing row updates, + * imposed by the underlying storage systems. The Trino engine classifies + * these different paradigms as elements of this RowChangeParadigm + * enumeration, returned by {@link ConnectorMetadata#getRowChangeParadigm} + */ +public enum RowChangeParadigm +{ + /** + * A storage paradigm in which the connector can update individual columns + * of rows identified by a rowId. The corresponding merge processor class is + * {@code ChangeOnlyUpdatedColumnsMergeProcessor} + */ + CHANGE_ONLY_UPDATED_COLUMNS, + + /** + * A paradigm that translates a changed row into a delete by rowId, and an insert of a + * new record, which will get a new rowId when the connector writes it out. The + * corresponding merge processor class is {@code DeleteAndInsertMergeProcessor}. + */ + DELETE_ROW_AND_INSERT_ROW, +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMergeSink.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMergeSink.java new file mode 100644 index 0000000000000..d22f0915baa66 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMergeSink.java @@ -0,0 +1,61 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi.connector.classloader; + +import com.facebook.presto.common.Page; +import com.facebook.presto.spi.ConnectorMergeSink; +import com.facebook.presto.spi.classloader.ThreadContextClassLoader; +import io.airlift.slice.Slice; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +import static java.util.Objects.requireNonNull; + +public class ClassLoaderSafeConnectorMergeSink + implements ConnectorMergeSink +{ + private final ConnectorMergeSink delegate; + private final ClassLoader classLoader; + + public ClassLoaderSafeConnectorMergeSink(ConnectorMergeSink delegate, ClassLoader classLoader) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + this.classLoader = requireNonNull(classLoader, "classLoader is null"); + } + + @Override + public void storeMergedRows(Page page) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + delegate.storeMergedRows(page); + } + } + + @Override + public CompletableFuture> finish() + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + return delegate.finish(); + } + } + + @Override + public void abort() + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + delegate.abort(); + } + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMetadata.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMetadata.java index 4d8d3b8bbe3d5..65420e7ee2795 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMetadata.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorMetadata.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorNewTableLayout; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorResolvedIndex; @@ -43,6 +44,7 @@ import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; import com.facebook.presto.spi.connector.ConnectorPartitioningMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; @@ -652,6 +654,48 @@ public void finishUpdate(ConnectorSession session, ConnectorTableHandle tableHan } } + @Override + public RowChangeParadigm getRowChangeParadigm(ConnectorSession session, ConnectorTableHandle tableHandle) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + return delegate.getRowChangeParadigm(session, tableHandle); + } + } + + /** + * Get the column handle that will generate row IDs for the merge operation. + * These IDs will be passed to the {@link com.facebook.presto.spi.ConnectorMergeSink#storeMergedRows} + * method of the {@link com.facebook.presto.spi.ConnectorMergeSink} that created them. + */ + public ColumnHandle getMergeRowIdColumnHandle(ConnectorSession session, ConnectorTableHandle tableHandle) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + return delegate.getMergeRowIdColumnHandle(session, tableHandle); + } + } + + @Override + public Optional getMergeUpdateLayout(ConnectorSession session, ConnectorTableHandle tableHandle) + { + return delegate.getMergeUpdateLayout(session, tableHandle); + } + + @Override + public ConnectorMergeTableHandle beginMerge(ConnectorSession session, ConnectorTableHandle tableHandle) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + return delegate.beginMerge(session, tableHandle); + } + } + + @Override + public void finishMerge(ConnectorSession session, ConnectorMergeTableHandle mergeTableHandle, Collection fragments, Collection computedStatistics) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + delegate.finishMerge(session, mergeTableHandle, fragments, computedStatistics); + } + } + @Override public boolean supportsMetadataDelete(ConnectorSession session, ConnectorTableHandle tableHandle, Optional tableLayoutHandle) { diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorPageSinkProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorPageSinkProvider.java index 761b6537e2a14..93316d8afb225 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorPageSinkProvider.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeConnectorPageSinkProvider.java @@ -15,6 +15,8 @@ import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorMergeSink; +import com.facebook.presto.spi.ConnectorMergeTableHandle; import com.facebook.presto.spi.ConnectorOutputTableHandle; import com.facebook.presto.spi.ConnectorPageSink; import com.facebook.presto.spi.ConnectorSession; @@ -60,4 +62,12 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa return new ClassLoaderSafeConnectorPageSink(delegate.createPageSink(transactionHandle, session, deleteTableHandle, pageSinkContext), classLoader); } } + + @Override + public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorMergeTableHandle mergeHandle) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + return new ClassLoaderSafeConnectorMergeSink(delegate.createMergeSink(transactionHandle, session, mergeHandle), classLoader); + } + } } diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java index 7878f1c542bab..4490990c827ea 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java @@ -26,6 +26,7 @@ import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import java.util.List; +import java.util.Optional; import java.util.function.ToIntFunction; import static java.util.Objects.requireNonNull; @@ -64,7 +65,7 @@ public List listPartitionHandles(ConnectorTransactionH } @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { return delegate.getBucketNodeMap(transactionHandle, session, partitioningHandle, sortedNodes); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java b/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java index c13fdcf5fda4a..08fb56d1a7587 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java @@ -13,13 +13,16 @@ */ package com.facebook.presto.spi.plan; +import com.facebook.presto.common.type.Type; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorTableMetadata; +import com.facebook.presto.spi.MergeHandle; import com.facebook.presto.spi.NewTableLayout; import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.SourceLocation; import com.facebook.presto.spi.TableHandle; +import com.facebook.presto.spi.connector.RowChangeParadigm; import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.relation.VariableReferenceExpression; import com.fasterxml.jackson.annotation.JsonCreator; @@ -527,4 +530,104 @@ public String toString() return handle.toString(); } } + + public static class MergeTarget + extends WriterTarget + { + private final TableHandle handle; + private final Optional mergeHandle; + private final SchemaTableName schemaTableName; + private final MergeParadigmAndTypes mergeParadigmAndTypes; + + @JsonCreator + public MergeTarget( + @JsonProperty("handle") TableHandle handle, + @JsonProperty("mergeHandle") Optional mergeHandle, + @JsonProperty("schemaTableName") SchemaTableName schemaTableName, + @JsonProperty("mergeParadigmAndTypes") MergeParadigmAndTypes mergeParadigmAndTypes) + { + this.handle = requireNonNull(handle, "handle is null"); + this.mergeHandle = requireNonNull(mergeHandle, "mergeHandle is null"); + this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null"); + this.mergeParadigmAndTypes = requireNonNull(mergeParadigmAndTypes, "mergeElements is null"); + } + + @JsonProperty + public TableHandle getHandle() + { + return handle; + } + + @JsonProperty + public Optional getMergeHandle() + { + return mergeHandle; + } + + @JsonProperty + public SchemaTableName getSchemaTableName() + { + return schemaTableName; + } + + @JsonProperty + public MergeParadigmAndTypes getMergeParadigmAndTypes() + { + return mergeParadigmAndTypes; + } + + @Override + public ConnectorId getConnectorId() + { + return handle.getConnectorId(); + } + + @Override + public Optional> getOutputColumns() + { + return Optional.empty(); + } + + @Override + public String toString() + { + return handle.toString(); + } + } + + public static class MergeParadigmAndTypes + { + private final RowChangeParadigm paradigm; + private final List columnTypes; + private final Type rowIdType; + + @JsonCreator + public MergeParadigmAndTypes( + @JsonProperty("paradigm") RowChangeParadigm paradigm, + @JsonProperty("columnTypes") List columnTypes, + @JsonProperty("rowIdType") Type rowIdType) + { + this.paradigm = requireNonNull(paradigm, "paradigm is null"); + this.columnTypes = requireNonNull(columnTypes, "columnTypes is null"); + this.rowIdType = requireNonNull(rowIdType, "rowIdType is null"); + } + + @JsonProperty + public RowChangeParadigm getParadigm() + { + return paradigm; + } + + @JsonProperty + public List getColumnTypes() + { + return columnTypes; + } + + @JsonProperty + public Type getRowIdType() + { + return rowIdType; + } + } } diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java index 8e97faedb2019..2f13f6295bb74 100644 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java +++ b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java @@ -25,6 +25,7 @@ import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.function.ToIntFunction; @@ -50,13 +51,13 @@ public TpcdsNodePartitioningProvider(NodeManager nodeManager, int splitsPerNode) } @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { Set nodes = nodeManager.getRequiredWorkerNodes(); checkState(!nodes.isEmpty(), "No TPCDS nodes available"); // Split the data using split and skew by the number of nodes available. - return createBucketNodeMap(toIntExact((long) nodes.size() * splitsPerNode)); + return Optional.of(createBucketNodeMap(toIntExact((long) nodes.size() * splitsPerNode))); } @Override diff --git a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java index d0213c9b359bb..98e455ab66da6 100644 --- a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java +++ b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java @@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableList; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.function.ToIntFunction; @@ -48,12 +49,12 @@ public TpchNodePartitioningProvider(NodeManager nodeManager, int splitsPerNode) } @Override - public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) + public Optional getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List sortedNodes) { Set nodes = nodeManager.getRequiredWorkerNodes(); // Split the data using split and skew by the number of nodes available. - return createBucketNodeMap(toIntExact((long) nodes.size() * splitsPerNode)); + return Optional.of(createBucketNodeMap(toIntExact((long) nodes.size() * splitsPerNode))); } @Override