Add cast to common type for join keys in RewriteJoinDict

Add simplified query test to reproduce issue

2 files changed, 47 insertions, 12 deletions

diff --git a/ydb/core/kqp/ut/kqp_join_ut.cpp b/ydb/core/kqp/ut/kqp_join_ut.cpp
index e5fd92a58f..99a320eb76 100644
--- a/ydb/core/kqp/ut/kqp_join_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_join_ut.cpp
@@ -1239,6 +1239,35 @@ Y_UNIT_TEST_SUITE(KqpJoin) {
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
         CompareYson(R"([[36u]])", FormatResultSetYson(result.GetResultSet(0)));
     }
+
+    Y_UNIT_TEST_NEW_ENGINE(JoinMismatchDictKeyTypes) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        UNIT_ASSERT(session.ExecuteSchemeQuery(R"(
+            CREATE TABLE `/Root/Join_Uint64_1` (
+                Key Uint64,
+                Value String,
+                PRIMARY KEY (Key)
+            );
+        )").GetValueSync().IsSuccess());
+
+        auto result = session.ExecuteDataQuery(Q1_(R"(
+            SELECT t1.Key
+            FROM
+                (SELECT 2 AS Key, COUNT(*) AS Cnt FROM `/Root/Join_Uint64_1`) AS t1
+            LEFT JOIN
+            AS_TABLE(
+                AsList(AsStruct(2u AS Key1, "Val1" AS Value))
+            ) AS t2
+            ON t1.Key = t2.Key1
+        )"), TTxControl::BeginTx().CommitTx()).GetValueSync();
+
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        CompareYson(R"([[2]])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
 }
 
 } // namespace NKqp
diff --git a/ydb/library/yql/dq/opt/dq_opt_peephole.cpp b/ydb/library/yql/dq/opt/dq_opt_peephole.cpp
index ea8864e0d7..fd545c0a3a 100644
--- a/ydb/library/yql/dq/opt/dq_opt_peephole.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_peephole.cpp
@@ -69,7 +69,7 @@ std::pair<TExprNode::TListType, TExprNode::TListType> JoinKeysToAtoms(TExprConte
 }
 
 TExprNode::TPtr BuildDictKeySelector(TExprContext& ctx, TPositionHandle pos, const TExprNode::TListType& keyAtoms,
-    const TTypeAnnotationNode::TListType& keyDryTypes, bool optional)
+    const TTypeAnnotationNode::TListType& keyDryTypes, bool needCast)
 {
     YQL_ENSURE(keyAtoms.size() == keyDryTypes.size());
 
@@ -89,7 +89,7 @@ TExprNode::TPtr BuildDictKeySelector(TExprContext& ctx, TPositionHandle pos, con
     }
 
     if (keysTuple.size() == 1) {
-        return optional
+        return needCast
             ? Build<TCoLambda>(ctx, pos)
                 .Args({keySelectorArg})
                 .Body<TCoStrictCast>()
@@ -106,7 +106,7 @@ TExprNode::TPtr BuildDictKeySelector(TExprContext& ctx, TPositionHandle pos, con
     }
 
     auto type = ctx.MakeType<TOptionalExprType>(ctx.MakeType<TTupleExprType>(keyDryTypes));
-    return optional
+    return needCast
         ? Build<TCoLambda>(ctx, pos)
             .Args({keySelectorArg})
             .Body<TCoStrictCast>()
@@ -174,7 +174,7 @@ TExprNode::TPtr AddConvertedKeys(TExprNode::TPtr list, TExprContext& ctx, TExprN
 }
 
 TExprNode::TListType OriginalJoinOutputMembers(const TDqPhyMapJoin& mapJoin, TExprContext& ctx) {
-    const auto origItemType = mapJoin.Ref().GetTypeAnn()->GetKind() == ETypeAnnotationKind::List ? 
+    const auto origItemType = mapJoin.Ref().GetTypeAnn()->GetKind() == ETypeAnnotationKind::List ?
         mapJoin.Ref().GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>() :
         mapJoin.Ref().GetTypeAnn()->Cast<TFlowExprType>()->GetItemType()->Cast<TStructExprType>();
     TExprNode::TListType structMembers;
@@ -194,7 +194,7 @@ TExprNode::TListType OriginalJoinOutputMembers(const TDqPhyMapJoin& mapJoin, TEx
  *  - Explicitly convert right input to the dict
  *  - Use quite pretty trick: do `MapJoinCore` in `FlatMap`-lambda
  *    (rely on the fact that there will be only one element in the `FlatMap`-stream)
- *  - Align key types using `StrictCast`, use internal columns to store converted left keys   
+ *  - Align key types using `StrictCast`, use internal columns to store converted left keys
  */
 TExprBase DqPeepholeRewriteMapJoin(const TExprBase& node, TExprContext& ctx) {
     if (!node.Maybe<TDqPhyMapJoin>()) {
@@ -465,19 +465,25 @@ NNodes::TExprBase DqPeepholeRewriteJoinDict(const NNodes::TExprBase& node, TExpr
     const auto* leftRowType = GetSeqItemType(joinDict.LeftInput().Ref().GetTypeAnn())->Cast<TStructExprType>();
     const auto* rightRowType = GetSeqItemType(joinDict.RightInput().Ref().GetTypeAnn())->Cast<TStructExprType>();
 
-    bool optKeyLeft = false, optKeyRight = false, badKey = false;
+    bool castKeyLeft = false, castKeyRight = false, badKey = false;
     TTypeAnnotationNode::TListType keyTypeItems;
     keyTypeItems.reserve(leftKeys.size());
     for (auto i = 0U; i < leftKeys.size(); ++i) {
+        bool optKeyLeft = false, optKeyRight = false;
         auto leftKeyType = leftRowType->FindItemType(leftKeys[i]->Content());
         auto rightKeyType = rightRowType->FindItemType(rightKeys[i]->Content());
-        keyTypeItems.emplace_back(CommonType<true>(node.Pos(), DryType(leftKeyType, optKeyLeft, ctx), DryType(rightKeyType, optKeyRight, ctx), ctx));
-        badKey = !keyTypeItems.back();
+        auto leftDryType = DryType(leftKeyType, optKeyLeft, ctx);
+        auto rightDryType = DryType(rightKeyType, optKeyRight, ctx);
+        auto commonType = CommonType<true>(node.Pos(), leftDryType, rightDryType, ctx);
+        badKey = !commonType;
         if (badKey) {
-            YQL_CLOG(DEBUG, CoreDq) << "Not comparable keys in join: " << leftKeys[i]->Content()
-                << "(" << *leftKeyType << ") vs " << rightKeys[i]->Content() << "(" << *rightKeyType << ")";
+            YQL_CLOG(DEBUG, CoreDq) << "Join has null result in key comparison: " << leftKeys[i]->Content()
+                << "(" << *leftKeyType << ") and " << rightKeys[i]->Content() << "(" << *rightKeyType << ")";
             break;
         }
+        castKeyLeft = (!IsSameAnnotation(*leftDryType, *commonType) || optKeyLeft);
+        castKeyRight = (!IsSameAnnotation(*rightDryType, *commonType) || optKeyRight);
+        keyTypeItems.emplace_back(commonType);
     }
 
     TExprNode::TPtr leftKeySelector;
@@ -498,8 +504,8 @@ NNodes::TExprBase DqPeepholeRewriteJoinDict(const NNodes::TExprBase& node, TExpr
                 .Build()
             .Done().Ptr();
     } else {
-        leftKeySelector = BuildDictKeySelector(ctx, joinDict.Pos(), leftKeys, keyTypeItems, optKeyLeft);
-        rightKeySelector = BuildDictKeySelector(ctx, joinDict.Pos(), rightKeys, keyTypeItems, optKeyRight);
+        leftKeySelector = BuildDictKeySelector(ctx, joinDict.Pos(), leftKeys, keyTypeItems, castKeyLeft);
+        rightKeySelector = BuildDictKeySelector(ctx, joinDict.Pos(), rightKeys, keyTypeItems, castKeyRight);
     }
 
     auto streamToDict = [&ctx](const TExprBase& input, const TExprNode::TPtr& keySelector) {