fix returning clause in upsert/insert on missing input columns (ydb-platform#16060)

Author: gridnevvvit

ydb/core/kqp/opt/kqp_opt_kql.cpp

@@ -15,11 +15,55 @@ using namespace NYql::NNodes;
 
 namespace {
 
+TVector<TString> GetMissingInputColumnsForReturning(
+    const TKiWriteTable& write, const TCoAtomList& inputColumns)
+{
+    auto returnColumns = write.ReturningColumns().Cast<TCoAtomList>();
+    if (returnColumns.Ref().ChildrenSize() == 0) {
+        return {};
+    }
+
+    THashSet<TStringBuf> currentInput;
+    for (const auto& name : inputColumns) {
+        currentInput.insert(name.Value());
+    }
+
+    TVector<TString> result;
+    for(const auto& returnCol : returnColumns) {
+        if (!currentInput.contains(returnCol.Value())) {
+            result.push_back(TString(returnCol.Value()));
+        }
+    }
+
+    return result;
+}
+
 // Replace absent input columns to NULL to perform REPLACE via UPSERT
-std::pair<TExprBase, TCoAtomList> CreateRowsToReplace(const TExprBase& input,
+std::pair<TExprBase, TCoAtomList> ExtendInputRowsWithAbsentNullColumns(const TKiWriteTable& write, const TExprBase& input,
     const TCoAtomList& inputColumns, const TKikimrTableDescription& tableDesc,
     TPositionHandle pos, TExprContext& ctx)
 {
+    TVector<TString> maybeMissingColumnsToReplace;
+    const auto op = GetTableOp(write);
+    if (op == TYdbOperation::Replace) {
+        for(const auto&[name, _]: tableDesc.Metadata->Columns) {
+            maybeMissingColumnsToReplace.push_back(name);
+        }
+    }
+
+    if (op == TYdbOperation::InsertAbort || op == TYdbOperation::InsertRevert || op == TYdbOperation::Upsert) {
+        maybeMissingColumnsToReplace = GetMissingInputColumnsForReturning(write, inputColumns);
+        if (maybeMissingColumnsToReplace.size() > 0) {
+            for(const auto& inputCol: inputColumns) {
+                maybeMissingColumnsToReplace.push_back(TString(inputCol.Value()));
+            }
+        }
+    }
+
+    if (maybeMissingColumnsToReplace.size() == 0) {
+        return {input, inputColumns};
+    }
+
     THashSet<TStringBuf> inputColumnsSet;
     for (const auto& name : inputColumns) {
         inputColumnsSet.insert(name.Value());
@@ -32,7 +76,7 @@ std::pair<TExprBase, TCoAtomList> CreateRowsToReplace(const TExprBase& input,
     TVector<TCoAtom> writeColumns;
     TVector<TExprBase> writeMembers;
 
-    for (const auto& [name, _] : tableDesc.Metadata->Columns) {
+    for (const auto& name : maybeMissingColumnsToReplace) {
         TMaybeNode<TExprBase> memberValue;
         if (tableDesc.GetKeyColumnIndex(name) || inputColumnsSet.contains(name)) {
             memberValue = Build<TCoMember>(ctx, pos)
@@ -214,21 +258,18 @@ TCoAtomList BuildUpsertInputColumns(const TCoAtomList& inputColumns,
 }
 
 std::pair<TExprBase, TCoAtomList> BuildWriteInput(const TKiWriteTable& write, const TKikimrTableDescription& table,
-    const TCoAtomList& inputColumns, const TCoAtomList& autoIncrement, const bool isSink,
+    const TCoAtomList& inputColumns, const TCoAtomList& autoIncrement, const bool /*isSink*/,
     TPositionHandle pos, TExprContext& ctx)
 {
     auto input = write.Input();
-    const bool isWriteReplace = (GetTableOp(write) == TYdbOperation::Replace) && !isSink;
 
     TCoAtomList inputCols = BuildUpsertInputColumns(inputColumns, autoIncrement, pos, ctx);
 
     if (autoIncrement.Ref().ChildrenSize() > 0) {
         input = BuildKqlSequencer(input, table, inputCols, autoIncrement, pos, ctx);
     }
 
-    if (isWriteReplace) {
-        std::tie(input, inputCols) = CreateRowsToReplace(input, inputCols, table, write.Pos(), ctx);
-    }
+    std::tie(input, inputCols) = ExtendInputRowsWithAbsentNullColumns(write, input, inputCols, table, write.Pos(), ctx);
 
     auto baseInput = Build<TKqpWriteConstraint>(ctx, pos)
         .Input(input)
@@ -238,15 +279,35 @@ std::pair<TExprBase, TCoAtomList> BuildWriteInput(const TKiWriteTable& write, co
     return {baseInput, inputCols};
 }
 
+
+TCoAtomList ExtendGenerateOnInsertColumnsList(const TKiWriteTable& write, TCoAtomList& generateColumnsIfInsert, const TCoAtomList& inputColumns, const TCoAtomList& autoincrement, TExprContext& ctx) {
+    auto inputCols = BuildUpsertInputColumns(inputColumns, autoincrement, write.Pos(), ctx);
+    auto maybeMissingColumnsToReplace = GetMissingInputColumnsForReturning(write, inputCols);
+    TVector<TExprNode::TPtr> result;
+    result.reserve(generateColumnsIfInsert.Ref().ChildrenSize() + maybeMissingColumnsToReplace.size());
+    for(const auto& item: generateColumnsIfInsert) {
+        result.push_back(item.Ptr());
+    }
+
+    for(auto& name: maybeMissingColumnsToReplace) {
+        auto atom = TCoAtom(ctx.NewAtom(write.Pos(), name));
+        result.push_back(atom.Ptr());
+    }
+
+    return Build<TCoAtomList>(ctx, write.Pos()).Add(result).Done();
+}
+
 TExprBase BuildUpsertTable(const TKiWriteTable& write, const TCoAtomList& inputColumns,
     const TCoAtomList& autoincrement, const bool isSink,
     const TKikimrTableDescription& table, TExprContext& ctx)
 {
     auto generateColumnsIfInsertNode = GetSetting(write.Settings().Ref(), "generate_columns_if_insert");
     YQL_ENSURE(generateColumnsIfInsertNode);
     TCoAtomList generateColumnsIfInsert = TCoNameValueTuple(generateColumnsIfInsertNode).Value().Cast<TCoAtomList>();
-
     auto settings = FilterSettings(write.Settings().Ref(), {"AllowInconsistentWrites"}, ctx);
+
+    generateColumnsIfInsert = ExtendGenerateOnInsertColumnsList(write, generateColumnsIfInsert, inputColumns, autoincrement, ctx);
+
     settings = AddSetting(*settings, write.Pos(), "Mode", Build<TCoAtom>(ctx, write.Pos()).Value("upsert").Done().Ptr(), ctx);
     const auto [input, columns] = BuildWriteInput(write, table, inputColumns, autoincrement, isSink, write.Pos(), ctx);
     if (generateColumnsIfInsert.Ref().ChildrenSize() > 0) {
@@ -282,6 +343,8 @@ TExprBase BuildUpsertTableWithIndex(const TKiWriteTable& write, const TCoAtomLis
     YQL_ENSURE(generateColumnsIfInsertNode);
     TCoAtomList generateColumnsIfInsert = TCoNameValueTuple(generateColumnsIfInsertNode).Value().Cast<TCoAtomList>();
 
+    generateColumnsIfInsert = ExtendGenerateOnInsertColumnsList(write, generateColumnsIfInsert, inputColumns, autoincrement, ctx);
+
     auto effect = Build<TKqlUpsertRowsIndex>(ctx, write.Pos())
         .Table(BuildTableMeta(table, write.Pos(), ctx))
         .Input(input.Ptr())

ydb/core/kqp/ut/opt/kqp_returning_ut.cpp

@@ -174,7 +174,7 @@ Y_UNIT_TEST(ReplaceSerial) {
 
 Y_UNIT_TEST(ReturningSerial) {
     NKikimrConfig::TAppConfig appConfig;
-    auto serverSettings = TKikimrSettings().SetAppConfig(appConfig);
+    auto serverSettings = TKikimrSettings().SetAppConfig(appConfig).SetWithSampleTables(false);
     TKikimrRunner kikimr(serverSettings);
 
     auto client = kikimr.GetTableClient();
@@ -198,6 +198,55 @@ Y_UNIT_TEST(ReturningSerial) {
     auto resultCreate = session.ExecuteSchemeQuery(queryCreate).GetValueSync();
     UNIT_ASSERT_C(resultCreate.IsSuccess(), resultCreate.GetIssues().ToString());
 
+   {
+        const auto query = Q_(R"(
+            --!syntax_v1
+            INSERT INTO ReturningTable (key) VALUES(20000) RETURNING *
+        )");
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
+        UNIT_ASSERT(result.IsSuccess());
+
+        CompareYson(R"([[20000;#]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
+
+   {
+        const auto query = Q_(R"(
+            --!syntax_v1
+            UPSERT INTO ReturningTable (key) VALUES(20000) RETURNING *
+        )");
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
+        UNIT_ASSERT(result.IsSuccess());
+
+        CompareYson(R"([[20000;#]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
+    {
+        const auto query = Q_(R"(
+            --!syntax_v1
+            UPSERT INTO ReturningTable (key, value) VALUES(20000, 100) RETURNING *
+        )");
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
+        UNIT_ASSERT(result.IsSuccess());
+
+        CompareYson(R"([[20000;[100]]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
+    {
+        const auto query = Q_(R"(
+            --!syntax_v1
+            UPSERT INTO ReturningTable (key) VALUES(20000) RETURNING *
+        )");
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).GetValueSync();
+        UNIT_ASSERT(result.IsSuccess());
+
+        CompareYson(R"([[20000;[100]]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
     {
         const auto query = Q_(R"(
             --!syntax_v1