Final changes for REPLACE/UPSERT/INSERT EXCLUDED (#831)

This commit includes the following: - Support for INSERT INTO tbl ... ON CONFLICT DO REPLACE EXCLUDED and REPLACE INTO tbl ... evaluation. - Support for INSERT INTO tbl ... ON CONFLICT DO UPDATE EXCLUDED logical planning. - The reason that evaluation is excluded is b/c the logic is for update evaluation requires implementing merge with existing values which makes it more completed; since it's not part of the related ask (attached issue) we defer the implementation for now - Adds parsing support for AS alias for UPSERT INTO and REPLACE INTO statements.
partiql · Oct 13, 2022 · 0b8d14a · 0b8d14a
1 parent 13f622f
commit 0b8d14a
Show file tree

Hide file tree

Showing 13 changed files with 286 additions and 20 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -42,9 +42,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
       [#707](https://github.com/partiql/partiql-lang-kotlin/issues/707), [#683](https://github.com/partiql/partiql-lang-kotlin/issues/683),
       and [#730](https://github.com/partiql/partiql-lang-kotlin/issues/730)
 - Parsing of `INSERT` DML with `ON CONFLICT DO REPLACE EXCLUDED` based on [RFC-0011](https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md)
-- Adds experimental parsing of `REPLACE INTO` and `UPSERT INTO` DML commands pending approval of the following RFC for moving out of experimental:
-  - https://github.com/partiql/partiql-docs/issues/27
-- Logical plan representation of `INSERT` DML with `ON CONFLICT DO REPLACE EXCLUDED` based on [RFC-0011](https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md)
+- Adds a subset of `REPLACE INTO` and `UPSERT INTO` parsing based on [RFC-0030](https://github.com/partiql/partiql-docs/blob/main/RFCs/0030-partiql-upsert-replace.md)
+  - Parsing of target attributes is not supported yet and is pending [#841](https://github.com/partiql/partiql-lang-kotlin/issues/841)
+- Logical plan representation and evaluation support for `INSERT` DML with `ON CONFLICT DO REPLACE EXCLUDED` and `REPLACE INTO` based on [RFC-0011](https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md)
+- Logical plan representation of `INSERT` DML with `ON CONFLICT DO UPDATE EXCLUDED` and `UPSERT INTO` based on [RFC-0011](https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md)
 - Enabled projection alias support for ORDER BY clause
 - Adds support for PIVOT in the planner consistent with `EvaluatingCompiler`
 

diff --git a/lang/resources/org/partiql/type-domains/partiql.ion b/lang/resources/org/partiql/type-domains/partiql.ion
@@ -732,8 +732,7 @@ may then be further optimized by selecting better implementations of each operat
                 (dml_insert)
                 (dml_delete)
                 (dml_replace)
-                // TODO:
-                // (dml_update)
+                (dml_update)
             )
         )
 

diff --git a/lang/src/org/partiql/lang/compiler/PartiQLCompilerDefault.kt b/lang/src/org/partiql/lang/compiler/PartiQLCompilerDefault.kt
@@ -97,6 +97,7 @@ internal class PartiQLCompilerDefault(
         return when (DmlAction.safeValueOf(action)) {
             DmlAction.INSERT -> PartiQLResult.Insert(target, rows)
             DmlAction.DELETE -> PartiQLResult.Delete(target, rows)
+            DmlAction.REPLACE -> PartiQLResult.Replace(target, rows)
             null -> error("Unknown DML Action `$action`")
         }
     }

diff --git a/lang/src/org/partiql/lang/eval/PartiQLResult.kt b/lang/src/org/partiql/lang/eval/PartiQLResult.kt
@@ -30,4 +30,9 @@ sealed class PartiQLResult {
         val target: String,
         val rows: Iterable<ExprValue>
     ) : PartiQLResult()
+
+    class Replace(
+        val target: String,
+        val rows: Iterable<ExprValue>
+    ) : PartiQLResult()
 }
diff --git a/lang/src/org/partiql/lang/planner/QueryPlan.kt b/lang/src/org/partiql/lang/planner/QueryPlan.kt
@@ -55,7 +55,8 @@ sealed class QueryResult {
  */
 enum class DmlAction {
     INSERT,
-    DELETE;
+    DELETE,
+    REPLACE;
 
     companion object {
         fun safeValueOf(v: String): DmlAction? = try {

diff --git a/lang/src/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransform.kt b/lang/src/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransform.kt
@@ -294,9 +294,19 @@ internal class AstToLogicalVisitorTransform(
                             } else -> TODO("Only `DO REPLACE EXCLUDED` is supported in logical plan at the moment.")
                         }
                     }
-                    is PartiqlAst.ConflictAction.DoUpdate -> TODO(
-                        "`ON CONFLICT DO UPDATE` is not supported in logical plan yet."
-                    )
+                    is PartiqlAst.ConflictAction.DoUpdate -> {
+                        when (conflictAction.value) {
+                            PartiqlAst.OnConflictValue.Excluded() -> PartiqlLogical.build {
+                                dml(
+                                    target = dmlOp.target.toDmlTargetId(),
+                                    operation = dmlUpdate(),
+                                    rows = transformExpr(dmlOp.values),
+                                    metas = node.metas
+                                )
+                            }
+                            else -> TODO("Only `DO UPDATE EXCLUDED` is supported in logical plan at the moment.")
+                        }
+                    }
                     is PartiqlAst.ConflictAction.DoNothing -> TODO(
                         "`ON CONFLICT DO NOTHING` is not supported in logical plan yet."
                     )

diff --git a/...c/org/partiql/lang/planner/transforms/LogicalResolvedToDefaultPhysicalVisitorTransform.kt b/...c/org/partiql/lang/planner/transforms/LogicalResolvedToDefaultPhysicalVisitorTransform.kt
@@ -153,8 +153,9 @@ internal class LogicalResolvedToDefaultPhysicalVisitorTransform(
         val action = when (node.operation) {
             is PartiqlLogicalResolved.DmlOperation.DmlInsert -> DmlAction.INSERT
             is PartiqlLogicalResolved.DmlOperation.DmlDelete -> DmlAction.DELETE
-            is PartiqlLogicalResolved.DmlOperation.DmlReplace ->
-                TODO("DmlReplace physical transform hasn't been implemented yet")
+            is PartiqlLogicalResolved.DmlOperation.DmlReplace -> DmlAction.REPLACE
+            is PartiqlLogicalResolved.DmlOperation.DmlUpdate ->
+                TODO("DmlUpdate physical transform is not supported yet")
         }.name.toLowerCase()
 
         return PartiqlPhysical.build {

diff --git a/lang/src/org/partiql/lang/visitors/PartiQLVisitor.kt b/lang/src/org/partiql/lang/visitors/PartiQLVisitor.kt
@@ -237,25 +237,33 @@ internal class PartiQLVisitor(val ion: IonSystem, val customTypes: List<CustomTy
 
     // TODO move from experimental; pending: https://github.com/partiql/partiql-docs/issues/27
     override fun visitReplaceCommand(ctx: PartiQLParser.ReplaceCommandContext) = PartiqlAst.build {
+        val asIdent = ctx.asIdent()
+        // Based on the RFC, if alias exists the table must be hidden behind the alias, see:
+        // https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md#41-insert-parameters
+        val target = if (asIdent != null) visitAsIdent(asIdent) else visitSymbolPrimitive(ctx.symbolPrimitive())
         insert(
-            target = visitSymbolPrimitive(ctx.symbolPrimitive()),
+            target = target,
             values = visit(ctx.value, PartiqlAst.Expr::class),
             conflictAction = doReplace(excluded()),
             metas = ctx.REPLACE().getSourceMetaContainer()
         )
     }
 
-    // TODO move from experimental; pending: https://github.com/partiql/partiql-docs/issues/27
+    // Based on https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md
     override fun visitUpsertCommand(ctx: PartiQLParser.UpsertCommandContext) = PartiqlAst.build {
+        val asIdent = ctx.asIdent()
+        // Based on the RFC, if alias exists the table must be hidden behind the alias, see:
+        // https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md#41-insert-parameters
+        val target = if (asIdent != null) visitAsIdent(asIdent) else visitSymbolPrimitive(ctx.symbolPrimitive())
         insert(
-            target = visitSymbolPrimitive(ctx.symbolPrimitive()),
+            target = target,
             values = visit(ctx.value, PartiqlAst.Expr::class),
             conflictAction = doUpdate(excluded()),
             metas = ctx.UPSERT().getSourceMetaContainer()
         )
     }
 
-    // FIXME: See `FIXME #001` in file `PartiQL.g4`.
+    // Based on https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md
     override fun visitInsertCommandReturning(ctx: PartiQLParser.InsertCommandReturningContext) = PartiqlAst.build {
         val metas = ctx.INSERT().getSourceMetaContainer()
         val target = visitPathSimple(ctx.pathSimple())
@@ -307,6 +315,7 @@ internal class PartiQLVisitor(val ion: IonSystem, val customTypes: List<CustomTy
         when {
             ctx.NOTHING() != null -> doNothing()
             ctx.REPLACE() != null -> visitDoReplace(ctx.doReplace())
+            ctx.UPDATE() != null -> visitDoUpdate(ctx.doUpdate())
             else -> TODO("ON CONFLICT only supports `DO REPLACE` and `DO NOTHING` actions at the moment.")
         }
     }
@@ -319,6 +328,14 @@ internal class PartiQLVisitor(val ion: IonSystem, val customTypes: List<CustomTy
         doReplace(value)
     }
 
+    override fun visitDoUpdate(ctx: PartiQLParser.DoUpdateContext) = PartiqlAst.build {
+        val value = when {
+            ctx.EXCLUDED() != null -> excluded()
+            else -> TODO("DO UPDATE doesn't support values other than `EXCLUDED` yet.")
+        }
+        doUpdate(value)
+    }
+
     override fun visitPathSimple(ctx: PartiQLParser.PathSimpleContext) = PartiqlAst.build {
         val root = visitSymbolPrimitive(ctx.symbolPrimitive())
         if (ctx.pathSimpleSteps().isEmpty()) return@build root

diff --git a/lang/test/org/partiql/lang/eval/evaluatortestframework/PartiQLCompilerPipelineFactory.kt b/lang/test/org/partiql/lang/eval/evaluatortestframework/PartiQLCompilerPipelineFactory.kt
@@ -93,7 +93,8 @@ internal class PartiQLCompilerPipelineFactory : PipelineFactory {
                 val statement = pipeline.compile(query)
                 return when (val result = statement.eval(session)) {
                     is PartiQLResult.Delete,
-                    is PartiQLResult.Insert -> error("DML is not supported by test suite")
+                    is PartiQLResult.Insert,
+                    is PartiQLResult.Replace -> error("DML is not supported by test suite")
                     is PartiQLResult.Value -> result.value
                 }
             }

diff --git a/lang/test/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransformTests.kt b/lang/test/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransformTests.kt
@@ -592,6 +592,68 @@ class AstToLogicalVisitorTransformTests {
                     }
                 }
             ),
+            TestCase(
+                "INSERT INTO foo SELECT x.* FROM 1 AS x ON CONFLICT DO UPDATE EXCLUDED",
+                PartiqlLogical.build {
+                    dml(
+                        identifier("foo", caseInsensitive()),
+                        dmlUpdate(),
+                        bindingsToValues(
+                            struct(structFields(id("x", caseInsensitive(), unqualified()))),
+                            scan(lit(ionInt(1)), varDecl("x"))
+                        )
+                    )
+                }
+            ),
+            TestCase(
+                "INSERT INTO foo AS f <<{'id': 1, 'name':'bob'}>> ON CONFLICT DO UPDATE EXCLUDED",
+                PartiqlLogical.build {
+                    PartiqlLogical.build {
+                        dml(
+                            identifier("f", caseInsensitive()),
+                            dmlUpdate(),
+                            bag(
+                                struct(
+                                    structField(lit(ionString("id")), lit(ionInt(1))),
+                                    structField(lit(ionString("name")), lit(ionString("bob")))
+                                )
+                            )
+                        )
+                    }
+                }
+            ),
+            TestCase(
+                "REPLACE INTO foo AS f <<{'id': 1, 'name':'bob'}>>",
+                PartiqlLogical.build {
+                    PartiqlLogical.build {
+                        dml(
+                            identifier("f", caseInsensitive()),
+                            dmlReplace(),
+                            bag(
+                                struct(
+                                    structField(lit(ionString("id")), lit(ionInt(1))),
+                                    structField(lit(ionString("name")), lit(ionString("bob")))
+                                )
+                            )
+                        )
+                    }
+                }
+            ),
+            TestCase(
+                "UPSERT INTO foo AS f SELECT x.* FROM 1 AS x",
+                PartiqlLogical.build {
+                    PartiqlLogical.build {
+                        dml(
+                            identifier("f", caseInsensitive()),
+                            dmlUpdate(),
+                            bindingsToValues(
+                                struct(structFields(id("x", caseInsensitive(), unqualified()))),
+                                scan(lit(ionInt(1)), varDecl("x"))
+                            )
+                        )
+                    }
+                }
+            ),
             TestCase(
                 "DELETE FROM y AS y",
                 PartiqlLogical.build {

diff --git a/.../partiql/lang/planner/transforms/LogicalResolvedToDefaultPhysicalVisitorTransformTests.kt b/.../partiql/lang/planner/transforms/LogicalResolvedToDefaultPhysicalVisitorTransformTests.kt
@@ -200,6 +200,38 @@ class LogicalResolvedToDefaultPhysicalVisitorTransformTests {
                     )
                 }
             ),
+            DmlTestCase(
+                // INSERT INTO foo SELECT x.* FROM 1 AS x ON CONFLICT DO REPLACE EXCLUDED
+                PartiqlLogicalResolved.build {
+                    dml(
+                        uniqueId = "foo",
+                        operation = dmlReplace(),
+                        rows = bindingsToValues(
+                            struct(structFields(localId(0))),
+                            scan(lit(ionInt(1)), varDecl(0))
+                        )
+                    )
+                },
+                PartiqlPhysical.build {
+                    dmlQuery(
+                        struct(
+                            structField(DML_COMMAND_FIELD_ACTION, "replace"),
+                            structField(DML_COMMAND_FIELD_TARGET_UNIQUE_ID, lit(ionSymbol("foo"))),
+                            structField(
+                                DML_COMMAND_FIELD_ROWS,
+                                bindingsToValues(
+                                    struct(structFields(localId(0))),
+                                    scan(
+                                        i = DEFAULT_IMPL,
+                                        expr = lit(ionInt(1)),
+                                        asDecl = varDecl(0)
+                                    )
+                                )
+                            )
+                        )
+                    )
+                }
+            ),
             DmlTestCase(
                 // DELETE FROM y AS y
                 PartiqlLogicalResolved.build {