Add SQL Support for MERGE INTO In Presto #20578 (engine-tests)

acarpente-denodo · David Stryker · acarpente-denodo · commit fdad084bc934 · 2025-11-19T11:51:24.000+01:00
Automated tests. Cherry-pick of trinodb/trino@cee96c3 Co-authored-by: David Stryker <david.stryker@starburstdata.com>
diff --git a/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.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;
@@ -35,6 +36,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;
@@ -460,6 +462,49 @@ public void finishUpdate(Session session, TableHandle tableHandle, Collection<Sl
         throw new UnsupportedOperationException();
     }
 
+    /**
+     * Return the row update paradigm supported by the connector on the table or throw
+     * an exception if row change is not supported.
+     */
+    public RowChangeParadigm getRowChangeParadigm(Session session, TableHandle tableHandle)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * 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.
+     */
+    public ColumnHandle getMergeTargetTableRowIdColumnHandle(Session session, TableHandle tableHandle)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Get the physical layout for updated or deleted rows of a MERGE operation.
+     */
+    public Optional<PartitioningHandle> getMergeUpdateLayout(Session session, TableHandle tableHandle)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Begin merge query
+     */
+    public MergeHandle beginMerge(Session session, TableHandle tableHandle)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Finish merge query
+     */
+    public void finishMerge(Session session, MergeHandle tableHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public Optional<ConnectorId> getCatalogHandle(Session session, String catalogName)
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
@@ -46,6 +46,7 @@
 import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorTableMetadata;
+import com.facebook.presto.spi.MaterializedViewDefinition;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.analyzer.ViewDefinition;
@@ -81,6 +82,7 @@
 import org.testng.annotations.BeforeClass;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Consumer;
@@ -337,6 +339,33 @@ public void setup()
                         ColumnMetadata.builder().setName("z").setType(BIGINT).build())),
                 false));
 
+        // materialized view referencing table in same schema
+        List<SchemaTableName> baseTables = new ArrayList<>(Collections.singletonList(table2));
+        MaterializedViewDefinition.TableColumn baseTableColumns = new MaterializedViewDefinition.TableColumn(table2, "a", true);
+
+        SchemaTableName materializedTable = new SchemaTableName("s1", "mv1");
+        MaterializedViewDefinition.TableColumn materializedViewTableColumn = new MaterializedViewDefinition.TableColumn(materializedTable, "a", true);
+
+        List<MaterializedViewDefinition.ColumnMapping> columnMappings = Collections.singletonList(
+                new MaterializedViewDefinition.ColumnMapping(materializedViewTableColumn, Collections.singletonList(baseTableColumns)));
+
+        MaterializedViewDefinition materializedViewData1 = new MaterializedViewDefinition(
+                        "select a from t2",
+                        "s1",
+                        "mv1",
+                        baseTables,
+                        Optional.of("user"),
+                        Optional.empty(),
+                        columnMappings,
+                        new ArrayList<>(),
+                        Optional.of(new ArrayList<>(Collections.singletonList("a"))));
+
+        ConnectorTableMetadata materializedViewMetadata1 = new ConnectorTableMetadata(
+                materializedTable, ImmutableList.of(ColumnMetadata.builder().setName("a").setType(BIGINT).build()));
+
+        inSetupTransaction(session ->
+                metadata.createMaterializedView(session, TPCH_CATALOG, materializedViewMetadata1, materializedViewData1, false));
+
         // valid view referencing table in same schema
         String viewData1 = JsonCodec.jsonCodec(ViewDefinition.class).toJson(
                 new ViewDefinition(
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
@@ -2351,4 +2351,17 @@ public void testTableFunctionRequiredColumns()
         // table s1.t5 has two columns. The second column is hidden. Table function can require a hidden column.
         analyze("SELECT * FROM TABLE(system.required_columns_function(input => TABLE(s1.t5)))");
     }
+
+    @Test
+    public void testInvalidMerge()
+    {
+        assertFails(MISSING_TABLE, "Table tpch.s1.foo does not exist",
+                "MERGE INTO foo USING bar ON foo.id = bar.id WHEN MATCHED THEN UPDATE SET id = bar.id + 1");
+
+        assertFails(NOT_SUPPORTED, "line 1:1: Merging into views is not supported",
+                "MERGE INTO v1 USING t1 ON v1.a = t1.a WHEN MATCHED THEN UPDATE SET id = bar.id + 1");
+
+        assertFails(NOT_SUPPORTED, "line 1:1: Merging into materialized views is not supported",
+                "MERGE INTO mv1 USING t1 ON mv1.a = t1.a WHEN MATCHED THEN  UPDATE SET id = bar.id + 1");
+    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestPruneMergeSourceColumns.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestPruneMergeSourceColumns.java
@@ -0,0 +1,79 @@
+/*
+ * 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.spi.SchemaTableName;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.sql.planner.assertions.PlanMatchPattern;
+import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
+import com.facebook.presto.sql.planner.plan.MergeWriterNode;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.List;
+
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.strictProject;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values;
+
+public class TestPruneMergeSourceColumns
+        extends BaseRuleTest
+{
+    @Test
+    public void testPruneInputColumn()
+    {
+        tester().assertThat(new PruneMergeSourceColumns())
+                .on(p -> {
+                    VariableReferenceExpression a = p.variable("a");
+                    VariableReferenceExpression mergeRow = p.variable("merge_row");
+                    VariableReferenceExpression rowId = p.variable("row_id");
+                    VariableReferenceExpression partialRows = p.variable("partial_rows");
+                    VariableReferenceExpression fragment = p.variable("fragment");
+                    List<VariableReferenceExpression> mergeProcessorProjectedSymbols = ImmutableList.of(mergeRow, rowId);
+                    return p.merge(
+                            new SchemaTableName("schema", "table"),
+                            p.values(a, mergeRow, rowId),
+                            mergeProcessorProjectedSymbols,
+                            ImmutableList.of(partialRows, fragment));
+                })
+                .matches(
+                        node(
+                                MergeWriterNode.class,
+                                strictProject(
+                                        ImmutableMap.of(
+                                                "row_id", PlanMatchPattern.expression("row_id"),
+                                                "merge_row", PlanMatchPattern.expression("merge_row")),
+                                        values("a", "merge_row", "row_id"))));
+    }
+
+    @Test
+    public void testDoNotPruneRowId()
+    {
+        tester().assertThat(new PruneMergeSourceColumns())
+                .on(p -> {
+                    VariableReferenceExpression mergeRow = p.variable("merge_row");
+                    VariableReferenceExpression rowId = p.variable("row_id");
+                    VariableReferenceExpression partialRows = p.variable("partial_rows");
+                    VariableReferenceExpression fragment = p.variable("fragment");
+                    List<VariableReferenceExpression> mergeProcessorProjectedSymbols = ImmutableList.of(mergeRow, rowId);
+                    return p.merge(
+                            new SchemaTableName("schema", "table"),
+                            p.values(mergeRow, rowId),
+                            mergeProcessorProjectedSymbols,
+                            ImmutableList.of(partialRows, fragment));
+                })
+                .doesNotFire();
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java
@@ -26,6 +26,7 @@
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.WarningCollector;
+import com.facebook.presto.spi.connector.RowChangeParadigm;
 import com.facebook.presto.spi.constraints.TableConstraint;
 import com.facebook.presto.spi.function.FunctionHandle;
 import com.facebook.presto.spi.plan.AggregationNode;
@@ -63,6 +64,8 @@
 import com.facebook.presto.spi.plan.TableFinishNode;
 import com.facebook.presto.spi.plan.TableScanNode;
 import com.facebook.presto.spi.plan.TableWriterNode;
+import com.facebook.presto.spi.plan.TableWriterNode.MergeParadigmAndTypes;
+import com.facebook.presto.spi.plan.TableWriterNode.MergeTarget;
 import com.facebook.presto.spi.plan.TopNNode;
 import com.facebook.presto.spi.plan.UnionNode;
 import com.facebook.presto.spi.plan.UnnestNode;
@@ -84,6 +87,7 @@
 import com.facebook.presto.sql.planner.plan.ExchangeNode;
 import com.facebook.presto.sql.planner.plan.GroupIdNode;
 import com.facebook.presto.sql.planner.plan.LateralJoinNode;
+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;
@@ -113,6 +117,7 @@
 
 import static com.facebook.presto.common.block.SortOrder.ASC_NULLS_FIRST;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
 import static com.facebook.presto.common.type.UnknownType.UNKNOWN;
 import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
 import static com.facebook.presto.spi.plan.ExchangeEncoding.COLUMNAR;
diff --git a/presto-spi/src/test/java/com/facebook/presto/spi/connector/classloader/TestClassLoaderSafeWrappers.java b/presto-spi/src/test/java/com/facebook/presto/spi/connector/classloader/TestClassLoaderSafeWrappers.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.spi.connector.classloader;
 
+import com.facebook.presto.spi.ConnectorMergeSink;
 import com.facebook.presto.spi.ConnectorPageSink;
 import com.facebook.presto.spi.connector.ConnectorMetadata;
 import com.facebook.presto.spi.connector.ConnectorNodePartitioningProvider;
@@ -34,5 +35,6 @@ public void testAllMethodsOverridden()
         assertAllMethodsOverridden(ConnectorPageSourceProvider.class, ClassLoaderSafeConnectorPageSourceProvider.class);
         assertAllMethodsOverridden(ConnectorSplitManager.class, ClassLoaderSafeConnectorSplitManager.class);
         assertAllMethodsOverridden(ConnectorNodePartitioningProvider.class, ClassLoaderSafeNodePartitioningProvider.class);
+        assertAllMethodsOverridden(ConnectorMergeSink.class, ClassLoaderSafeConnectorMergeSink.class);
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@`
`13`	`13`	`*/`
`14`	`14`	`package com.facebook.presto.spi.connector.classloader;`
`15`	`15`
	`16`	`+import com.facebook.presto.spi.ConnectorMergeSink;`
`16`	`17`	`import com.facebook.presto.spi.ConnectorPageSink;`
`17`	`18`	`import com.facebook.presto.spi.connector.ConnectorMetadata;`
`18`	`19`	`import com.facebook.presto.spi.connector.ConnectorNodePartitioningProvider;`
`@@ -34,5 +35,6 @@ public void testAllMethodsOverridden()`
`34`	`35`	`assertAllMethodsOverridden(ConnectorPageSourceProvider.class, ClassLoaderSafeConnectorPageSourceProvider.class);`
`35`	`36`	`assertAllMethodsOverridden(ConnectorSplitManager.class, ClassLoaderSafeConnectorSplitManager.class);`
`36`	`37`	`assertAllMethodsOverridden(ConnectorNodePartitioningProvider.class, ClassLoaderSafeNodePartitioningProvider.class);`
	`38`	`+ assertAllMethodsOverridden(ConnectorMergeSink.class, ClassLoaderSafeConnectorMergeSink.class);`
`37`	`39`	`}`
`38`	`40`	`}`