Allow UPDATE ON, UPSERT with partial set of input... (#894)

* Allow UPDATE ON, UPSERT with partial set of input...

columns in case of unique index.

Example:
  table: pk, fk1, fk2. Uniq index: fk1, fk2.

To perform UPSERT INTO table (pk, fk1) we
need to read missed value fk2 from main table to
create lookup uniq index key.

Right now we have two lookups in to main table
- first one has been described above
- second we need to make to remove old value
from index. It will be fixed in a next step.

* fix

9 files changed, 1105 insertions, 71 deletions

diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects.cpp b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects.cpp
index 54d2e81e05..119bf07332 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects.cpp
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects.cpp
@@ -77,7 +77,6 @@ TMaybe<TCondenseInputResult> CondenseInput(const TExprBase& input, TExprContext&
     TVector<TCoArgument> stageArguments;
 
     if (IsDqPureExpr(input)) {
-        YQL_ENSURE(input.Ref().GetTypeAnn()->GetKind() == ETypeAnnotationKind::List, "" << input.Ref().Dump());
         auto stream = Build<TCoToStream>(ctx, input.Pos())
             .Input<TCoJust>()
                 .Input(input)
diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects_impl.h b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects_impl.h
index 257e389368..e2be88de4c 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects_impl.h
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_effects_impl.h
@@ -32,9 +32,17 @@ TMaybe<TCondenseInputResult> CondenseInputToDictByPk(const NYql::NNodes::TExprBa
 
 NYql::NNodes::TMaybeNode<NYql::NNodes::TDqPhyPrecompute> PrecomputeTableLookupDict(
     const NYql::NNodes::TDqPhyPrecompute& lookupKeys, const NYql::TKikimrTableDescription& table,
+    const TVector<NYql::NNodes::TExprBase>& columnsList,
+    NYql::TPositionHandle pos, NYql::TExprContext& ctx, bool fixLookupKeys);
+
+NYql::NNodes::TMaybeNode<NYql::NNodes::TDqPhyPrecompute> PrecomputeTableLookupDict(
+    const NYql::NNodes::TDqPhyPrecompute& lookupKeys, const NYql::TKikimrTableDescription& table,
     const THashSet<TString>& dataColumns, const THashSet<TString>& keyColumns, NYql::TPositionHandle pos,
     NYql::TExprContext& ctx);
 
+NYql::NNodes::TDqPhyPrecompute PrecomputeCondenseInputResult(const TCondenseInputResult& condenseResult,
+    NYql::TPositionHandle pos, NYql::TExprContext& ctx);
+
 // Creates key selector using PK of given table
 NYql::NNodes::TCoLambda MakeTableKeySelector(const NYql::TKikimrTableMetadataPtr tableMeta, NYql::TPositionHandle pos,
     NYql::TExprContext& ctx, TMaybe<int> tupleId = {});
diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_indexes.cpp b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_indexes.cpp
index c9232fd7f2..3bb2e9a50d 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_indexes.cpp
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_indexes.cpp
@@ -48,7 +48,11 @@ TVector<TExprBase> CreateColumnsToSelectToUpdateIndex(
     return columnsToSelect;
 }
 
-TDqPhyPrecompute PrecomputeDict(const TCondenseInputResult& condenseResult, TPositionHandle pos, TExprContext& ctx) {
+} // namespace
+
+TDqPhyPrecompute PrecomputeCondenseInputResult(const TCondenseInputResult& condenseResult,
+    TPositionHandle pos, TExprContext& ctx)
+{
     auto computeDictStage = Build<TDqStage>(ctx, pos)
         .Inputs()
             .Add(condenseResult.StageInputs)
@@ -70,8 +74,6 @@ TDqPhyPrecompute PrecomputeDict(const TCondenseInputResult& condenseResult, TPos
         .Done();
 }
 
-} // namespace
-
 TVector<std::pair<TExprNode::TPtr, const TIndexDescription*>> BuildSecondaryIndexVector(
     const TKikimrTableDescription& table,
     TPositionHandle pos,
@@ -127,26 +129,81 @@ TSecondaryIndexes BuildSecondaryIndexVector(const TKikimrTableDescription& table
 }
 
 TMaybeNode<TDqPhyPrecompute> PrecomputeTableLookupDict(const TDqPhyPrecompute& lookupKeys,
-    const TKikimrTableDescription& table, const THashSet<TString>& dataColumns,
-    const THashSet<TString>& keyColumns, TPositionHandle pos, TExprContext& ctx)
+    const TKikimrTableDescription& table, const TVector<TExprBase>& columnsList,
+    TPositionHandle pos, TExprContext& ctx, bool fixLookupKeys)
 {
-    auto lookupColumns = CreateColumnsToSelectToUpdateIndex(table.Metadata->KeyColumnNames, dataColumns,
-        keyColumns, pos, ctx);
-
     auto lookupColumnsList = Build<TCoAtomList>(ctx, pos)
-        .Add(lookupColumns)
+        .Add(columnsList)
         .Done();
 
+    TExprNode::TPtr keys;
+
+    // we need to left only table key columns to perform lookup
+    // unfortunately we can't do it inside lookup stage
+    if (fixLookupKeys) {
+        auto keyArg = TCoArgument(ctx.NewArgument(pos, "key"));
+        auto keysList = TCoArgument(ctx.NewArgument(pos, "keys_list"));
+        TVector<TExprBase> keyLookupTuples;
+        keyLookupTuples.reserve(table.Metadata->KeyColumnNames.size());
+
+        for (const auto& key : table.Metadata->KeyColumnNames) {
+            keyLookupTuples.emplace_back(
+                Build<TCoNameValueTuple>(ctx, pos)
+                    .Name().Build(key)
+                    .Value<TCoMember>()
+                        .Struct(keyArg)
+                        .Name().Build(key)
+                        .Build()
+                    .Done());
+        }
+
+        auto list = Build<TCoToStream>(ctx, pos)
+            .Input<TCoJust>()
+                .Input<TCoMap>()
+                    .Input(keysList)
+                    .Lambda()
+                        .Args({keyArg})
+                        .Body<TCoAsStruct>()
+                            .Add(keyLookupTuples)
+                            .Build()
+                        .Build()
+                    .Build()
+                .Build()
+            .Done().Ptr();
+
+        keys = Build<TDqStage>(ctx, pos)
+            .Inputs()
+                .Add(lookupKeys)
+                .Build()
+            .Program()
+                .Args({keysList})
+                .Body(list)
+                .Build()
+            .Settings().Build()
+            .Done().Ptr();
+
+        keys = Build<TDqPhyPrecompute>(ctx, pos)
+            .Connection<TDqCnValue>()
+                .Output()
+                    .Stage(keys)
+                    .Index().Build("0")
+                    .Build()
+                .Build()
+            .Done().Ptr();
+    } else {
+        keys = lookupKeys.Ptr();
+    }
+
     auto lookupStage = Build<TDqStage>(ctx, pos)
         .Inputs()
-            .Add(lookupKeys)
+            .Add(keys)
             .Build()
         .Program()
-            .Args({"keys_list"})
+            .Args({"keys_stage_arg"})
             .Body<TKqpLookupTable>()
                 .Table(BuildTableMeta(table, pos, ctx))
                 .LookupKeys<TCoIterator>()
-                    .List("keys_list")
+                    .List("keys_stage_arg")
                     .Build()
                 .Columns(lookupColumnsList)
                 .Build()
@@ -167,7 +224,17 @@ TMaybeNode<TDqPhyPrecompute> PrecomputeTableLookupDict(const TDqPhyPrecompute& l
         return {};
     }
 
-    return PrecomputeDict(*condenseLookupResult, lookupKeys.Pos(), ctx);
+    return PrecomputeCondenseInputResult(*condenseLookupResult, lookupKeys.Pos(), ctx);
+}
+
+TMaybeNode<TDqPhyPrecompute> PrecomputeTableLookupDict(const TDqPhyPrecompute& lookupKeys,
+    const TKikimrTableDescription& table, const THashSet<TString>& dataColumns,
+    const THashSet<TString>& keyColumns, TPositionHandle pos, TExprContext& ctx)
+{
+    auto lookupColumns = CreateColumnsToSelectToUpdateIndex(table.Metadata->KeyColumnNames, dataColumns,
+        keyColumns, pos, ctx);
+
+    return PrecomputeTableLookupDict(lookupKeys, table, lookupColumns, pos, ctx, false);
 }
 
 TExprBase MakeRowsFromDict(const TDqPhyPrecompute& dict, const TVector<TString>& dictKeys,
diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.cpp b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.cpp
index 811aad6e40..31dcdbd928 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.cpp
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.cpp
@@ -264,17 +264,17 @@ TVector<TUniqBuildHelper::TUniqCheckNodes> TUniqBuildHelper::Prepare(const TCoAr
         // Compatibility with PG semantic - allow multiple null in columns with unique constaint
         TVector<TCoAtom> skipNullColumns;
         skipNullColumns.reserve(table.Metadata->Indexes[i].KeyColumns.size());
+
+        bool used = false;
         for (const auto& column : table.Metadata->Indexes[i].KeyColumns) {
-            if (!inputColumns || inputColumns->contains(column)) {
-                TCoAtom atom(ctx.NewAtom(pos, column));
-                skipNullColumns.emplace_back(atom);
-            }
+            used |= (!inputColumns || inputColumns->contains(column));
+            TCoAtom atom(ctx.NewAtom(pos, column));
+            skipNullColumns.emplace_back(atom);
         }
 
-        //no columns to skip -> no index columns to check -> skip check
-        if (skipNullColumns.empty()) {
-            continue;
-        }
+        // Just to doublecheck we are not trying to update index without data to update
+        YQL_ENSURE(used, "Index is used but not input columns for update. Probably it's a bug."
+            " Index: " << table.Metadata->Indexes[i].Name);
 
         auto skipNull = Build<TCoSkipNullMembers>(ctx, pos)
             .Input(rowsListArg)
@@ -292,11 +292,46 @@ TVector<TUniqBuildHelper::TUniqCheckNodes> TUniqBuildHelper::Prepare(const TCoAr
     return checks;
 }
 
+static TExprNode::TPtr CreateRowsToPass(const TCoArgument& rowsListArg, const THashSet<TStringBuf>* inputColumns,
+    TPositionHandle pos, TExprContext& ctx)
+{
+    if (!inputColumns) {
+        return rowsListArg.Ptr();
+    }
+
+    auto arg = TCoArgument(ctx.NewArgument(pos, "arg"));
+
+    TVector<TExprBase> columns;
+    columns.reserve(inputColumns->size());
+
+    for (const auto x : *inputColumns) {
+        columns.emplace_back(
+            Build<TCoNameValueTuple>(ctx, pos)
+                .Name().Build(x)
+                .Value<TCoMember>()
+                    .Struct(arg)
+                    .Name().Build(x)
+                    .Build()
+                .Done());
+    }
+
+    return Build<TCoMap>(ctx, pos)
+        .Input(rowsListArg)
+        .Lambda()
+            .Args({arg})
+            .Body<TCoAsStruct>()
+                .Add(columns)
+                .Build()
+            .Build()
+        .Done().Ptr();
+}
+
 TUniqBuildHelper::TUniqBuildHelper(const TKikimrTableDescription& table, const THashSet<TStringBuf>* inputColumns, const THashSet<TString>* usedIndexes,
     TPositionHandle pos, TExprContext& ctx, bool skipPkCheck)
     : RowsListArg(ctx.NewArgument(pos, "rows_list"))
     , False(MakeBool(pos, false, ctx))
     , Checks(Prepare(RowsListArg, table, inputColumns, usedIndexes, pos, ctx, skipPkCheck))
+    , RowsToPass(CreateRowsToPass(RowsListArg, inputColumns, pos, ctx))
 {}
 
 TUniqBuildHelper::TUniqCheckNodes TUniqBuildHelper::MakeUniqCheckNodes(const TCoLambda& selector,
@@ -340,7 +375,7 @@ TDqStage TUniqBuildHelper::CreateComputeKeysStage(const TCondenseInputResult& co
 
     types.emplace_back(
         Build<TCoTypeOf>(ctx, pos)
-            .Value(RowsListArg)
+            .Value(RowsToPass)
             .Done()
     );
 
@@ -376,7 +411,7 @@ TDqStage TUniqBuildHelper::CreateComputeKeysStage(const TCondenseInputResult& co
 
     variants.emplace_back(
         Build<TCoVariant>(ctx, pos)
-            .Item(RowsListArg)
+            .Item(RowsToPass)
             .Index().Build("0")
             .VarType(variantType)
             .Done()
diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.h b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.h
index 65c5d07678..285dfca576 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.h
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_uniq_helper.h
@@ -76,6 +76,7 @@ protected:
     const NYql::TExprNode::TPtr False;
 private:
     const TChecks Checks;
+    const NYql::TExprNode::TPtr RowsToPass;
 };
 
 TUniqBuildHelper::TPtr CreateInsertUniqBuildHelper(const NYql::TKikimrTableDescription& table,
diff --git a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp
index 8295db03dd..405c9e0e3f 100644
--- a/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp
+++ b/ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp
@@ -10,6 +10,8 @@ using namespace NYql;
 using namespace NYql::NDq;
 using namespace NYql::NNodes;
 
+//#define OPT_IDX_DEBUG 1
+
 namespace {
 
 struct TRowsAndKeysResult {
@@ -342,30 +344,166 @@ TExprBase MakeUpsertIndexRows(TKqpPhyUpsertIndexMode mode, const TDqPhyPrecomput
         .Done();
 }
 
-TMaybe<TCondenseInputResult> CheckUniqueConstraint(const TExprBase& inputRows, const THashSet<TStringBuf> inputColumns, bool checkOnlyGivenColumns,
-    const TKikimrTableDescription& table, const TSecondaryIndexes& indexes, TPositionHandle pos, TExprContext& ctx)
+TMaybe<TCondenseInputResult> CheckUniqueConstraint(const TExprBase& inputRows, const THashSet<TStringBuf> inputColumns,
+    const TKikimrTableDescription& table, const TSecondaryIndexes& indexes,
+    TPositionHandle pos, TExprContext& ctx)
 {
     auto condenseResult = CondenseInput(inputRows, ctx);
     if (!condenseResult) {
         return {};
     }
 
+    // In case of absent on of index columns in the UPSERT(or UPDATE ON) values
+    // we need to get actual value from main table to perform lookup in to the index table
+    THashSet<std::string_view> missedKeyInput;
+
     // Check uniq constraint for indexes which will be updated by input data.
     // but skip main table pk columns - handle case where we have a complex index is a tuple contains pk
     const auto& mainPk = table.Metadata->KeyColumnNames;
     THashSet<TString> usedIndexes;
+    bool hasUniqIndex = false;
     for (const auto& [_, indexDesc] : indexes) {
+        hasUniqIndex |= (indexDesc->Type == TIndexDescription::EType::GlobalSyncUnique);
         for (const auto& indexKeyCol : indexDesc->KeyColumns) {
-            if (inputColumns.contains(indexKeyCol) &&
-                (std::find(mainPk.begin(), mainPk.end(), indexKeyCol) == mainPk.end()))
+            if (inputColumns.contains(indexKeyCol)
+                && std::find(mainPk.begin(), mainPk.end(), indexKeyCol) == mainPk.end())
             {
                 usedIndexes.insert(indexDesc->Name);
-                break;
+            } else {
+                // input always contains key columns
+                YQL_ENSURE(std::find(mainPk.begin(), mainPk.end(), indexKeyCol) == mainPk.end());
+                missedKeyInput.emplace(indexKeyCol);
             }
         }
     }
 
-    auto helper = CreateUpsertUniqBuildHelper(table, checkOnlyGivenColumns ? &inputColumns : nullptr, usedIndexes, pos, ctx);
+    if (missedKeyInput && hasUniqIndex) {
+        TVector<TExprBase> columns;
+
+        TCoArgument inLambdaArg(ctx.NewArgument(pos, "in_lambda_arg"));
+        auto missedFromMain = TCoArgument(ctx.NewArgument(pos, "missed_from_main"));
+
+        TVector<TExprBase> resCol;
+
+        for (const auto& x : inputColumns) {
+            if (!missedKeyInput.contains(x)) {
+                resCol.emplace_back(
+                    Build<TCoNameValueTuple>(ctx, pos)
+                        .Name().Build(x)
+                        .Value<TCoMember>()
+                            .Struct(inLambdaArg)
+                            .Name().Build(x)
+                            .Build()
+                        .Done());
+            }
+        }
+
+        TVector<TExprBase> resNullCol = resCol;
+
+        for (const auto& x : missedKeyInput) {
+            auto atom = Build<TCoAtom>(ctx, pos)
+                .Value(x)
+                .Done();
+            columns.emplace_back(atom);
+
+            auto columnType = table.GetColumnType(TString(x));
+            YQL_ENSURE(columnType);
+
+            resCol.emplace_back(
+                    Build<TCoNameValueTuple>(ctx, pos)
+                        .Name().Build(x)
+                        .Value<TCoMember>()
+                            .Struct(missedFromMain)
+                            .Name().Build(x)
+                            .Build()
+                        .Done());
+
+            resNullCol.emplace_back(
+                Build<TCoNameValueTuple>(ctx, pos)
+                    .Name().Build(x)
+                    .Value<TCoNothing>()
+                        .OptionalType(NCommon::BuildTypeExpr(pos, *columnType, ctx))
+                        .Build()
+                    .Done());
+        }
+
+        for (const auto& x : mainPk) {
+            auto atom = Build<TCoAtom>(ctx, pos)
+                .Value(x)
+                .Done();
+            columns.emplace_back(atom);
+        }
+
+        auto inPrecompute = PrecomputeCondenseInputResult(*condenseResult, pos, ctx);
+
+        auto precomputeTableLookupDict = PrecomputeTableLookupDict(inPrecompute, table, columns, pos, ctx, true);
+
+        TVector<TExprBase> keyLookupTuples;
+        for (const auto& key : mainPk) {
+            keyLookupTuples.emplace_back(
+                Build<TCoNameValueTuple>(ctx, pos)
+                    .Name().Build(key)
+                    .Value<TCoMember>()
+                        .Struct(inLambdaArg)
+                        .Name().Build(key)
+                        .Build()
+                    .Done());
+        }
+
+        TCoArgument inPrecomputeArg(ctx.NewArgument(pos, "in_precompute_arg"));
+        TCoArgument lookupPrecomputeArg(ctx.NewArgument(pos, "lookup_precompute_arg"));
+        auto fillMissedStage = Build<TDqStage>(ctx, pos)
+            .Inputs()
+                .Add(inPrecompute)
+                .Add(precomputeTableLookupDict.Cast())
+                .Build()
+            .Program()
+                .Args({inPrecomputeArg, lookupPrecomputeArg})
+                .Body<TCoToStream>()
+                    .Input<TCoJust>()
+                        .Input<TCoMap>()
+                            .Input(inPrecomputeArg)
+                            .Lambda()
+                                .Args({inLambdaArg})
+                                .Body<TCoIfPresent>()
+                                    .Optional<TCoLookup>()
+                                        .Collection(lookupPrecomputeArg)
+                                        .Lookup<TCoAsStruct>()
+                                            .Add(keyLookupTuples)
+                                            .Build()
+                                        .Build()
+                                    .PresentHandler<TCoLambda>()
+                                        .Args({missedFromMain})
+                                        .Body<TCoAsStruct>()
+                                            .Add(resCol)
+                                            .Build()
+                                        .Build()
+                                    .MissingValue<TCoAsStruct>()
+                                        .Add(resNullCol)
+                                        .Build()
+                                    .Build()
+                                .Build()
+                            .Build()
+                        .Build()
+                    .Build()
+                .Build()
+            .Settings().Build()
+            .Done();
+
+        auto connection = Build<TDqPhyPrecompute>(ctx, pos)
+            .Connection<TDqCnValue>()
+                .Output()
+                    .Stage(fillMissedStage)
+                    .Index().Build("0")
+                    .Build()
+                .Build()
+            .Done();
+
+        condenseResult = CondenseInput(connection, ctx);
+        YQL_ENSURE(condenseResult);
+    }
+
+    auto helper = CreateUpsertUniqBuildHelper(table, &inputColumns, usedIndexes, pos, ctx);
     if (helper->GetChecksNum() == 0) {
         return condenseResult;
     }
@@ -474,20 +612,10 @@ TMaybeNode<TExprList> KqpPhyUpsertIndexEffectsImpl(TKqpPhyUpsertIndexMode mode,
         inputColumnsSet.emplace(column.Value());
     }
 
-    bool checkOnlyGivenColumns = false;
-    if (settings) {
-        for (const auto& setting : settings.Cast()) {
-            if (setting.Name().Value() == "IsUpdate") {
-                checkOnlyGivenColumns = true;
-                break;
-            }
-        }
-    }
-
     auto filter =  (mode == TKqpPhyUpsertIndexMode::UpdateOn) ? &inputColumnsSet : nullptr;
     const auto indexes = BuildSecondaryIndexVector(table, pos, ctx, filter);
 
-    auto checkedInput = CheckUniqueConstraint(inputRows, inputColumnsSet, checkOnlyGivenColumns, table, indexes, pos, ctx);
+    auto checkedInput = CheckUniqueConstraint(inputRows, inputColumnsSet, table, indexes, pos, ctx);
 
     if (!checkedInput) {
         return {};
@@ -751,9 +879,15 @@ TMaybeNode<TExprList> KqpPhyUpsertIndexEffectsImpl(TKqpPhyUpsertIndexMode mode,
         }
     }
 
-    return Build<TExprList>(ctx, pos)
+    auto ret = Build<TExprList>(ctx, pos)
         .Add(effects)
         .Done();
+
+#ifdef OPT_IDX_DEBUG
+    Cerr << KqpExprToPrettyString(ret, ctx) << Endl;
+#endif
+
+    return ret;
 }
 
 TExprBase KqpBuildUpsertIndexStages(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx) {
diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
index 62d751870f..5bce05bb68 100644
--- a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
+++ b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
@@ -366,7 +366,7 @@ void FillEffectRows(const TEffectCallable& callable, TEffectProto& proto, bool i
     }
 }
 
-void FillLookup(const TKqpLookupTable& lookup, NKqpProto::TKqpPhyOpLookup& lookupProto) {
+void FillLookup(const TKqpLookupTable& lookup, NKqpProto::TKqpPhyOpLookup& lookupProto, TExprContext& ctx) {
     auto maybeList = lookup.LookupKeys().Maybe<TCoIterator>().List();
     YQL_ENSURE(maybeList, "Expected iterator as lookup input, got: " << lookup.LookupKeys().Ref().Content());
 
@@ -398,7 +398,9 @@ void FillLookup(const TKqpLookupTable& lookup, NKqpProto::TKqpPhyOpLookup& looku
             }
         }
     } else {
-        YQL_ENSURE(false, "Unexpected lookup input: " << maybeList.Cast().Ref().Content());
+        auto brokenLookup =  KqpExprToPrettyString(lookup, ctx);
+        YQL_ENSURE(false, "Unexpected lookup input: " << maybeList.Cast().Ref().Content()
+            << "lookup: " << brokenLookup);
     }
 }
 
@@ -621,7 +623,7 @@ private:
                 FillTablesMap(lookupTable.Table(), lookupTable.Columns(), tablesMap);
                 FillTableId(lookupTable.Table(), *tableOp.MutableTable());
                 FillColumns(lookupTable.Columns(), *tableMeta, tableOp, true);
-                FillLookup(lookupTable, *tableOp.MutableLookup());
+                FillLookup(lookupTable, *tableOp.MutableLookup(), ctx);
             } else if (auto maybeUpsertRows = node.Maybe<TKqpUpsertRows>()) {
                 auto upsertRows = maybeUpsertRows.Cast();
                 auto tableMeta = TablesData->ExistingTable(Cluster, upsertRows.Table().Path()).Metadata;
diff --git a/ydb/core/kqp/ut/indexes/kqp_indexes_multishard_ut.cpp b/ydb/core/kqp/ut/indexes/kqp_indexes_multishard_ut.cpp
index 57944ffa03..4367907ca9 100644
--- a/ydb/core/kqp/ut/indexes/kqp_indexes_multishard_ut.cpp
+++ b/ydb/core/kqp/ut/indexes/kqp_indexes_multishard_ut.cpp
@@ -283,7 +283,7 @@ Y_UNIT_TEST_SUITE(KqpUniqueIndex) {
             const TString expected = R"([[[1000000000u];[1u]];[[1000000001u];[2u]];[[3000000000u];[3u]];[[4294967295u];[4u]]])";
             UNIT_ASSERT_VALUES_EQUAL(yson, expected);
         }
-return;
+
         {
             const TString query(Q_(R"(
                 UPDATE `/Root/MultiShardIndexed` SET fk = 1000000000 WHERE value = "v1";
@@ -495,11 +495,86 @@ return;
             auto result = ExecuteDataQuery(session, query);
             UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
         }
+
         {
             const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed/index/indexImplTable");
             const TString expected = R"([[[1000000000u];[1u]];[[1000000000u];[2u]];[[3000000000u];[3u]];[[4294967295u];[4u]]])";
             UNIT_ASSERT_VALUES_EQUAL(yson, expected);
         }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed");
+            const TString expected = R"([[[1u];[1000000000u];["v1"]];[[2u];[1000000000u];["v2"]];[[3u];[3000000000u];["v3"]];[[4u];[4294967295u];["v4"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        // No such key - do nothing for update
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexed` ON (key, fk) VALUES (NULL, 1000000000u);
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed/index/indexImplTable");
+            const TString expected = R"([[[1000000000u];[1u]];[[1000000000u];[2u]];[[3000000000u];[3u]];[[4294967295u];[4u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed");
+            const TString expected = R"([[[1u];[1000000000u];["v1"]];[[2u];[1000000000u];["v2"]];[[3u];[3000000000u];["v3"]];[[4u];[4294967295u];["v4"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        // Check correct handle only one NULL in pk
+        for (int i = 0; i < 2; i++) {
+            {
+                const TString query(Q_(R"(
+                    UPSERT INTO `/Root/MultiShardIndexed` (key, fk) VALUES (NULL, 1000000000u);
+                )"));
+
+                auto result = ExecuteDataQuery(session, query);
+                UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+            }
+
+            {
+                const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed/index/indexImplTable");
+                const TString expected = R"([[[1000000000u];#];[[1000000000u];[1u]];[[1000000000u];[2u]];[[3000000000u];[3u]];[[4294967295u];[4u]]])";
+                UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+            }
+
+            {
+                const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed");
+                const TString expected = R"([[#;[1000000000u];#];[[1u];[1000000000u];["v1"]];[[2u];[1000000000u];["v2"]];[[3u];[3000000000u];["v3"]];[[4u];[4294967295u];["v4"]]])";
+                UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+            }
+        }
+
+        // There is NULL key - update
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexed` ON (key, fk) VALUES (NULL, 90000000u);
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed/index/indexImplTable");
+            const TString expected = R"([[[90000000u];#];[[1000000000u];[1u]];[[1000000000u];[2u]];[[3000000000u];[3u]];[[4294967295u];[4u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexed");
+            const TString expected = R"([[#;[90000000u];#];[[1u];[1000000000u];["v1"]];[[2u];[1000000000u];["v2"]];[[3u];[3000000000u];["v3"]];[[4u];[4294967295u];["v4"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
     }
 
     Y_UNIT_TEST(InsertNullInPk) {
@@ -617,6 +692,455 @@ return;
         }
     }
 
+    Y_UNIT_TEST(UpsertExplicitNullInComplexFk) {
+        TKikimrRunner kikimr(SyntaxV1Settings());
+        CreateTableWithMultishardIndexComplexFk(kikimr.GetTestClient(), IG_UNIQUE);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, fk2, value) VALUES
+                (1173915, NULL, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, fk2, value) VALUES
+                (1173916, NULL, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, fk2, value) VALUES
+                (1173917, 1, NULL, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]];[[1u];#;[1173917u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+    }
+
+    Y_UNIT_TEST(UpsertImplicitNullInComplexFk) {
+        TKikimrRunner kikimr(SyntaxV1Settings());
+        CreateTableWithMultishardIndexComplexFk(kikimr.GetTestClient(), IG_UNIQUE);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk2, value) VALUES
+                (1173915, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk2, value) VALUES
+                (1173916, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, value) VALUES
+                (1173917, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]];[[1u];#;[1173917u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk2, value) VALUES
+                (1173917, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]];[[1u];[1u];[1173917u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, value) VALUES
+                (1173916, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]];[[1u];[1u];[1173917u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+    }
+
+    Y_UNIT_TEST(UpdateImplicitNullInComplexFk2) {
+        TKikimrRunner kikimr(SyntaxV1Settings());
+        CreateTableWithMultishardIndexComplexFk(kikimr.GetTestClient(), IG_UNIQUE);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, fk2, value) VALUES
+                (1173915, NULL, 1, "v1"),
+                (1173916, 1, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[[1u];[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, value) VALUES
+                (1173915, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[[1u];[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` SET fk1 = 1 WHERE key = 1173915;
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[[1u];[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+    }
+
+    Y_UNIT_TEST(UpdateOnNullInComplexFk) {
+        TKikimrRunner kikimr(SyntaxV1Settings());
+        CreateTableWithMultishardIndexComplexFk(kikimr.GetTestClient(), IG_UNIQUE);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk2, value) VALUES
+                (1173915, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk2, value) VALUES
+                (1173916, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, value) VALUES
+                (1173915, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[1u];[1u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        // we need new row
+        {
+            const TString query(Q_(R"(
+                UPSERT INTO `/Root/MultiShardIndexedComplexFk` (key, fk1, fk2, value) VALUES
+                (1173916, 2, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[1u];[1u];[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173915, 2, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[1u];[1u];[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk2, value) VALUES
+                (1173915, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[1u];[2u];[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173915, NULL, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173916, NULL, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173915u]];[#;[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173915, NULL, 1, "v1"),
+                (1173916, NULL, 1, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[1u];[1173915u]];[#;[1u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk2, value) VALUES
+                (1173915, 2, "v1"),
+                (1173916, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173915u]];[#;[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173915, 2, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173916u]];[[2u];[2u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, value) VALUES
+                (1173916, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173916u]];[[2u];[2u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173916, 2, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[#;[2u];[1173916u]];[[2u];[2u];[1173915u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk1, fk2, value) VALUES
+                (1173915, 2, NULL, "v1"),
+                (1173916, 2, 2,    "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[2u];#;[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query(Q_(R"(
+                UPDATE `/Root/MultiShardIndexedComplexFk` ON (key, fk2, value) VALUES
+                (1173915, 2, "v1");
+            )"));
+
+            auto result = ExecuteDataQuery(session, query);
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto yson = ReadTableToYson(session, "/Root/MultiShardIndexedComplexFk/index/indexImplTable");
+            const TString expected = R"([[[2u];#;[1173915u]];[[2u];[2u];[1173916u]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+    }
+
     Y_UNIT_TEST(InsertNullInComplexFkDuplicate) {
         TKikimrRunner kikimr(SyntaxV1Settings());
         CreateTableWithMultishardIndexComplexFk(kikimr.GetTestClient(), IG_UNIQUE);
diff --git a/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp b/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp
index 929ca9d29f..a7feb778cf 100644
--- a/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp
+++ b/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp
@@ -594,6 +594,270 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
         UNIT_ASSERT_VALUES_EQUAL_C(result2.GetStatus(), NYdb::EStatus::ABORTED, result2.GetIssues().ToString().c_str());
     }
 
+    Y_UNIT_TEST(UniqIndexComplexPkComplexFkOverlap) {
+        auto setting = NKikimrKqp::TKqpSetting();
+        auto serverSettings = TKikimrSettings()
+            .SetKqpSettings({setting});
+        TKikimrRunner kikimr(serverSettings);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            auto tableBuilder = db.GetTableBuilder();
+            tableBuilder
+                .AddNullableColumn("k1", EPrimitiveType::String)
+                .AddNullableColumn("k2", EPrimitiveType::String)
+                .AddNullableColumn("fk1", EPrimitiveType::String)
+                .AddNullableColumn("fk2", EPrimitiveType::Int32)
+                .AddNullableColumn("fk3", EPrimitiveType::Uint64)
+                .AddNullableColumn("Value", EPrimitiveType::String);
+            tableBuilder.SetPrimaryKeyColumns(TVector<TString>{"k1", "k2"});
+            tableBuilder.AddUniqueSecondaryIndex("Index", TVector<TString>{"fk1", "fk2", "fk3", "k2"});
+            auto result = session.CreateTable("/Root/TestTable", tableBuilder.Build()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+        }
+
+        {
+            // Upsert - add new rows
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, k2, fk1, fk2, fk3, Value) VALUES
+                ("p1str1", NULL, "fk1_str", 1000, 1000000000u, "Value1_1"),
+                ("p1str2", NULL, "fk1_str", 1000, 1000000000u, "Value1_1");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
+            const TString expected = R"([[["fk1_str"];[1000];[1000000000u];#;["p1str1"]];[["fk1_str"];[1000];[1000000000u];#;["p1str2"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            // Upsert - add new rows
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, k2, fk1, fk2, fk3, Value) VALUES
+                ("p1str3", "p2str3", "fk1_str", 1000, 1000000000u, "Value1_1");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
+            const TString expected = R"([[["fk1_str"];[1000];[1000000000u];#;["p1str1"]];[["fk1_str"];[1000];[1000000000u];#;["p1str2"]];[["fk1_str"];[1000];[1000000000u];["p2str3"];["p1str3"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str3"];["p2str3"];["fk1_str"];[1000];[1000000000u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            // Upsert - add new pk but broke uniq constraint for index
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, k2, fk1, fk2, fk3, Value) VALUES
+                ("p1str4", "p2str3", "fk1_str", 1000, 1000000000u, "Value1_1");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            // Upsert - add new rows with NULL in one of fk
+            // set k2 to p2str3 to make conflict on the next itteration
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, k2, fk1, fk2, fk3, Value) VALUES
+                ("p1str5", "p2str3", NULL, 1000, 1000000000u, "Value1_1");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
+            const TString expected = R"([[#;[1000];[1000000000u];["p2str3"];["p1str5"]];[["fk1_str"];[1000];[1000000000u];#;["p1str1"]];[["fk1_str"];[1000];[1000000000u];#;["p1str2"]];[["fk1_str"];[1000];[1000000000u];["p2str3"];["p1str3"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str3"];["p2str3"];["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str5"];["p2str3"];#;[1000];[1000000000u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            // Upsert - conflict with [["p1str1"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]] during update [["p1str5"];["p2str3"];#;[1000];[1000000000u];["Value1_1"]]
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, k2, fk1) VALUES
+                ("p1str5", "p2str3", "fk1_str");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
+            const TString expected = R"([[#;[1000];[1000000000u];["p2str3"];["p1str5"]];[["fk1_str"];[1000];[1000000000u];#;["p1str1"]];[["fk1_str"];[1000];[1000000000u];#;["p1str2"]];[["fk1_str"];[1000];[1000000000u];["p2str3"];["p1str3"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];#;["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str3"];["p2str3"];["fk1_str"];[1000];[1000000000u];["Value1_1"]];[["p1str5"];["p2str3"];#;[1000];[1000000000u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+    }
+
+    Y_UNIT_TEST(UpsertMultipleUniqIndexes) {
+        auto setting = NKikimrKqp::TKqpSetting();
+        auto serverSettings = TKikimrSettings()
+            .SetKqpSettings({setting});
+        TKikimrRunner kikimr(serverSettings);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        {
+            auto tableBuilder = db.GetTableBuilder();
+            tableBuilder
+                .AddNullableColumn("k1", EPrimitiveType::String)
+                .AddNullableColumn("fk1", EPrimitiveType::String)
+                .AddNullableColumn("fk2", EPrimitiveType::Int32)
+                .AddNullableColumn("fk3", EPrimitiveType::Uint64)
+                .AddNullableColumn("Value", EPrimitiveType::String);
+            tableBuilder.SetPrimaryKeyColumns(TVector<TString>{"k1"});
+            tableBuilder.AddUniqueSecondaryIndex("Index12", TVector<TString>{"fk1", "fk2"}, TVector<TString>{"Value"});
+            tableBuilder.AddUniqueSecondaryIndex("Index3", TVector<TString>{"fk3"});
+            auto result = session.CreateTable("/Root/TestTable", tableBuilder.Build()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+        }
+
+        {
+            // Upsert - add new rows
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, fk1, fk2, fk3, Value) VALUES
+                ("p1str1", "fk1_str1", 1000, 1000000000u, "Value1_1"),
+                ("p1str2", "fk1_str2", 1000, 1000000001u, "Value1_1");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index12/indexImplTable");
+            const TString expected = R"([[["fk1_str1"];[1000];["p1str1"];["Value1_1"]];[["fk1_str2"];[1000];["p1str2"];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index3/indexImplTable");
+            const TString expected = R"([[[1000000000u];["p1str1"]];[[1000000001u];["p1str2"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];["fk1_str1"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];["fk1_str2"];[1000];[1000000001u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, fk3, Value) VALUES
+                ("p1str2", 1000000000u, "Value1_2");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index12/indexImplTable");
+            const TString expected = R"([[["fk1_str1"];[1000];["p1str1"];["Value1_1"]];[["fk1_str2"];[1000];["p1str2"];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index3/indexImplTable");
+            const TString expected = R"([[[1000000000u];["p1str1"]];[[1000000001u];["p1str2"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];["fk1_str1"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];["fk1_str2"];[1000];[1000000001u];["Value1_1"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const TString query2 = Q1_(R"(
+                UPSERT INTO `/Root/TestTable` (k1, fk3, Value) VALUES
+                ("p1str2", 2000000000u, "Value1_3");
+            )");
+
+            auto result = session.ExecuteDataQuery(
+                                 query2,
+                                 TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+                          .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index12/indexImplTable");
+            const TString expected = R"([[["fk1_str1"];[1000];["p1str1"];["Value1_1"]];[["fk1_str2"];[1000];["p1str2"];["Value1_3"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index3/indexImplTable");
+            const TString expected = R"([[[1000000000u];["p1str1"]];[[2000000000u];["p1str2"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+
+        {
+            const auto& yson = ReadTablePartToYson(session, "/Root/TestTable");
+            const TString expected = R"([[["p1str1"];["fk1_str1"];[1000];[1000000000u];["Value1_1"]];[["p1str2"];["fk1_str2"];[1000];[2000000000u];["Value1_3"]]])";
+            UNIT_ASSERT_VALUES_EQUAL(yson, expected);
+        }
+    }
+
     void DoUpsertWithoutIndexUpdate(bool uniq) {
         auto setting = NKikimrKqp::TKqpSetting();
         auto serverSettings = TKikimrSettings()
@@ -620,6 +884,7 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
             UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
         }
 
+        auto uniqExtraStages = uniq ? 6 : 0;
         {
             // Upsert - add new row
             const TString query2 = Q1_(R"(
@@ -636,26 +901,27 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
                           .ExtractValueSync();
             UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
             auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), 5);
+
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(),  uniqExtraStages + 5);
 
             // One read from main table
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access().size(), 1);
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access(0).name(), "/Root/TestTable");
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access(0).reads().rows(), 0);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access().size(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access(0).name(), "/Root/TestTable");
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access(0).reads().rows(), 0);
 
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(2).table_access().size(), 0);
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(3).table_access().size(), 0);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 2).table_access().size(), 0);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 3).table_access().size(), 0);
 
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access().size(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access().size(), 2);
 
             // One update of main table
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(0).name(), "/Root/TestTable");
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(0).updates().rows(), 1);
-            UNIT_ASSERT(            !stats.query_phases(4).table_access(0).has_deletes());
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(0).name(), "/Root/TestTable");
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(0).updates().rows(), 1);
+            UNIT_ASSERT(            !stats.query_phases(uniqExtraStages + 4).table_access(0).has_deletes());
 
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(1).name(), "/Root/TestTable/Index/indexImplTable");
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(1).updates().rows(), 1);
-            UNIT_ASSERT(            !stats.query_phases(4).table_access(1).has_deletes());
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(1).name(), "/Root/TestTable/Index/indexImplTable");
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(1).updates().rows(), 1);
+            UNIT_ASSERT(            !stats.query_phases(uniqExtraStages + 4).table_access(1).has_deletes());
 
             {
                 const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
@@ -680,21 +946,21 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
                           .ExtractValueSync();
             UNIT_ASSERT(result.IsSuccess());
             auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), 5);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), uniqExtraStages + 5);
 
             // One read from main table
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access().size(), 1);
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access(0).name(), "/Root/TestTable");
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(1).table_access(0).reads().rows(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access().size(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access(0).name(), "/Root/TestTable");
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 1).table_access(0).reads().rows(), 1);
 
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(2).table_access().size(), 0);
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(3).table_access().size(), 0);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 2).table_access().size(), 0);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 3).table_access().size(), 0);
 
             // One update of main table
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access().size(), 1);
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(0).name(), "/Root/TestTable");
-            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(4).table_access(0).updates().rows(), 1);
-            UNIT_ASSERT(            !stats.query_phases(4).table_access(0).has_deletes());
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access().size(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(0).name(), "/Root/TestTable");
+            UNIT_ASSERT_VALUES_EQUAL(stats.query_phases(uniqExtraStages + 4).table_access(0).updates().rows(), 1);
+            UNIT_ASSERT(            !stats.query_phases(uniqExtraStages + 4).table_access(0).has_deletes());
 
             {
                 const auto& yson = ReadTablePartToYson(session, "/Root/TestTable/Index/indexImplTable");
@@ -705,8 +971,6 @@ Y_UNIT_TEST_SUITE(KqpIndexes) {
     }
 
     Y_UNIT_TEST_TWIN(DoUpsertWithoutIndexUpdate, UniqIndex) {
-        if (UniqIndex)
-            return; // TODO: fixit!!!
         DoUpsertWithoutIndexUpdate(UniqIndex);
     }