Make key type non optional for inner/semi join.

3 files changed, 28 insertions, 17 deletions

diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index 5fb7b2a30c..86062fbe0b 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -1625,10 +1625,11 @@ const TTypeAnnotationNode* DryType(const TTypeAnnotationNode* type, TExprContext
     return nullptr;
 }
 
-const TTypeAnnotationNode* JoinDryKeyType(const TTypeAnnotationNode* primary, const TTypeAnnotationNode* secondary, TExprContext& ctx) {
-    if (const auto dry = DryType(primary, ctx))
+const TTypeAnnotationNode* JoinDryKeyType(bool outer, const TTypeAnnotationNode* primary, const TTypeAnnotationNode* secondary, TExprContext& ctx) {
+    bool hasOptional = false;
+    if (const auto dry = DryType(primary, hasOptional, ctx))
         if (!((NUdf::ECastOptions::AnywayLoseData | NUdf::ECastOptions::Impossible) & CastResult<true>(secondary, dry)))
-            return dry;
+            return hasOptional && outer ? ctx.MakeType<TOptionalExprType>(dry) : dry;
     return nullptr;
 }
 
@@ -1639,17 +1640,18 @@ const TTypeAnnotationNode* JoinDryKeyType(const TTypeAnnotationNode* primary, co
     return nullptr;
 }
 
-const TTypeAnnotationNode* JoinCommonDryType(TPositionHandle position, const TTypeAnnotationNode* one, const TTypeAnnotationNode* two, TExprContext& ctx) {
-    auto dryOne = DryType(one, ctx);
-    auto dryTwo = DryType(two, ctx);
+const TTypeAnnotationNode* JoinCommonDryKeyType(TPositionHandle position, bool outer, const TTypeAnnotationNode* one, const TTypeAnnotationNode* two, TExprContext& ctx) {
+    bool optOne = false, optTwo = false;
+    auto dryOne = DryType(one, optOne, ctx);
+    auto dryTwo = DryType(two, optTwo, ctx);
     if (!(dryOne && dryTwo))
         return nullptr;
 
-    if (dryOne->GetKind() == ETypeAnnotationKind::Optional && dryTwo->GetKind() != ETypeAnnotationKind::Optional)
-        dryTwo = ctx.MakeType<TOptionalExprType>(dryTwo);
-    else  if (dryOne->GetKind() != ETypeAnnotationKind::Optional && dryTwo->GetKind() == ETypeAnnotationKind::Optional)
+    if (outer && (optOne || optTwo)) {
         dryOne = ctx.MakeType<TOptionalExprType>(dryOne);
-
+        dryTwo = ctx.MakeType<TOptionalExprType>(dryTwo);
+    }
+    
     return CommonType<true>(position, dryOne, dryTwo, ctx);
 }
 
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index 6d32887125..59a2ecaa35 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -191,10 +191,10 @@ const TTypeAnnotationNode* DryType(const TTypeAnnotationNode* type, TExprContext
 const TTypeAnnotationNode* DryType(const TTypeAnnotationNode* type, bool& hasOptional, TExprContext& ctx);
 
 // Key type for left or right join.
-const TTypeAnnotationNode* JoinDryKeyType(const TTypeAnnotationNode* primary, const TTypeAnnotationNode* secondary, TExprContext& ctx);
+const TTypeAnnotationNode* JoinDryKeyType(bool outer, const TTypeAnnotationNode* primary, const TTypeAnnotationNode* secondary, TExprContext& ctx);
 const TTypeAnnotationNode* JoinDryKeyType(const TTypeAnnotationNode* primary, const TTypeAnnotationNode* secondary, bool& hasOptional, TExprContext& ctx);
 // Key type for inner or full join.
-const TTypeAnnotationNode* JoinCommonDryType(TPositionHandle position, const TTypeAnnotationNode* one, const TTypeAnnotationNode* two, TExprContext& ctx);
+const TTypeAnnotationNode* JoinCommonDryKeyType(TPositionHandle position, bool outer, const TTypeAnnotationNode* one, const TTypeAnnotationNode* two, TExprContext& ctx);
 
 template <bool Strict> // Strict + DryType before - common type for join key.
 const TTypeAnnotationNode* CommonType(TPositionHandle position, const TTypeAnnotationNode* one, const TTypeAnnotationNode* two, TExprContext& ctx);
diff --git a/ydb/library/yql/dq/opt/dq_opt_join.cpp b/ydb/library/yql/dq/opt/dq_opt_join.cpp
index 9ff01dcec4..5e99d0a8ad 100644
--- a/ydb/library/yql/dq/opt/dq_opt_join.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_join.cpp
@@ -951,14 +951,23 @@ bool PrepareJoinSide(
         .Seal().Build();
 
     if (filter) {
-        TExprNode::TListType check;
+        TExprNode::TListType check, unwrap;
         check.reserve(keys.size() + remap.size());
+        unwrap.reserve(remap.size());
         std::transform(keys.cbegin(), keys.cend(), std::back_inserter(check), [&](const TCoAtom& key) { return key.Ptr(); });
-        std::transform(remap.cbegin(), remap.cend(), std::back_inserter(check), [&](const TModifyKeysList::value_type& key) { return std::get<1>(key).Ptr(); });
+        std::for_each(remap.cbegin(), remap.cend(), [&](const TModifyKeysList::value_type& key) {
+            (ETypeAnnotationKind::Optional == std::get<const TTypeAnnotationNode*>(key)->GetKind() ? check : unwrap).emplace_back(std::get<1>(key).Ptr());
+        });
         preprocess = Build<TCoSkipNullMembers>(ctx, preprocess->Pos())
             .Input(std::move(preprocess))
             .Members().Add(std::move(check)).Build()
             .Done().Ptr();
+        if (!unwrap.empty()) {
+            preprocess = Build<TCoFilterNullMembers>(ctx, preprocess->Pos())
+                .Input(std::move(preprocess))
+                .Members().Add(std::move(unwrap)).Build()
+                .Done().Ptr();
+        }
     }
 
     const auto lambda = Build<TCoLambda>(ctx, preprocess->Pos())
@@ -1047,11 +1056,11 @@ TExprBase DqBuildHashJoin(const TDqJoin& join, EHashJoinMode mode, TExprContext&
         const auto keyType2 = rightStructType->FindItemType(rightJoinKeys[i]);
         const TTypeAnnotationNode* commonType = nullptr;
         if (leftKind) {
-            commonType = JoinDryKeyType(keyType1, keyType2, ctx);
+            commonType = JoinDryKeyType(!filter, keyType1, keyType2, ctx);
         } else if (rightKind){
-            commonType = JoinDryKeyType(keyType2, keyType1, ctx);
+            commonType = JoinDryKeyType(!filter, keyType2, keyType1, ctx);
         } else {
-            commonType = JoinCommonDryType(join.Pos(), keyType1, keyType2, ctx);
+            commonType = JoinCommonDryKeyType(join.Pos(), !filter, keyType1, keyType2, ctx);
         }
 
         if (commonType) {