allow JOINs for tables with not null pk

fix(kqp): allow join for tables with not null pk

2 files changed, 120 insertions, 2 deletions

diff --git a/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp b/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
index 19dc6a0ba2f..912dfca53c3 100644
--- a/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
@@ -12,8 +12,10 @@ bool GetEquiJoinKeyTypes(TExprBase leftInput, const TString& leftColumnName,
 {
     auto rightType = rightTable.GetColumnType(rightColumnName);
     YQL_ENSURE(rightType);
-    YQL_ENSURE(rightType->GetKind() == ETypeAnnotationKind::Optional);
-    rightType = rightType->Cast<TOptionalExprType>()->GetItemType();
+    if (rightType->GetKind() == ETypeAnnotationKind::Optional) {
+        rightType = rightType->Cast<TOptionalExprType>()->GetItemType();
+    }
+
     YQL_ENSURE(rightType->GetKind() == ETypeAnnotationKind::Data);
     rightData = rightType->Cast<TDataExprType>();
 
diff --git a/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp b/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
index 3ee6b55ad35..7a03f49a796 100644
--- a/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
@@ -715,6 +715,122 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
                 result.GetIssues().ToString());
         }
     }
+
+    Y_UNIT_TEST_NEW_ENGINE(JoinBothTablesWithNotNullPk) {
+        TKikimrRunner kikimr;
+        auto client = kikimr.GetTableClient();
+        auto session = client.CreateSession().GetValueSync().GetSession();
+
+        {
+            auto createTableResult = session.ExecuteSchemeQuery(Q1_(R"(
+                CREATE TABLE `/Root/Left` (
+                    Key Uint64 NOT NULL,
+                    Value String,
+                    PRIMARY KEY (Key)
+                );
+            )")).ExtractValueSync();
+            UNIT_ASSERT_C(createTableResult.IsSuccess(), createTableResult.GetIssues().ToString());
+
+            auto result = session.ExecuteDataQuery(Q1_(R"(
+                UPSERT INTO `/Root/Left` (Key, Value) VALUES (1, 'lValue1'), (2, 'lValue2');
+            )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            auto createTableResult = session.ExecuteSchemeQuery(Q1_(R"(
+                CREATE TABLE `/Root/Right` (
+                    Key Uint64 NOT NULL,
+                    Value String,
+                    PRIMARY KEY (Key)
+                );
+            )")).ExtractValueSync();
+            UNIT_ASSERT_C(createTableResult.IsSuccess(), createTableResult.GetIssues().ToString());
+
+            auto result = session.ExecuteDataQuery(Q1_(R"(
+                UPSERT INTO `/Root/Right` (Key, Value) VALUES (1, 'rValue1'), (3, 'rValue3');
+            )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const auto query = Q1_(R"(
+                SELECT l.Value, r.Value FROM `/Root/Left` AS l JOIN `/Root/Right` AS r ON l.Key = r.Key;
+            )");
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            CompareYson(R"([[["lValue1"];["rValue1"]]])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+    }
+
+    Y_UNIT_TEST_NEW_ENGINE(JoinLeftTableWithNotNullPk) {
+        TKikimrRunner kikimr;
+        auto client = kikimr.GetTableClient();
+        auto session = client.CreateSession().GetValueSync().GetSession();
+
+        {
+            auto createTableResult = session.ExecuteSchemeQuery(Q1_(R"(
+                CREATE TABLE `/Root/Left` (
+                    Key Uint64 NOT NULL,
+                    Value String,
+                    PRIMARY KEY (Key)
+                );
+            )")).ExtractValueSync();
+            UNIT_ASSERT_C(createTableResult.IsSuccess(), createTableResult.GetIssues().ToString());
+
+            auto result = session.ExecuteDataQuery(Q1_(R"(
+                UPSERT INTO `/Root/Left` (Key, Value) VALUES (1, 'lValue1'), (2, 'lValue2');
+            )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            auto createTableResult = session.ExecuteSchemeQuery(Q1_(R"(
+                CREATE TABLE `/Root/Right` (
+                    Key Uint64,
+                    Value String,
+                    PRIMARY KEY (Key)
+                );
+            )")).ExtractValueSync();
+            UNIT_ASSERT_C(createTableResult.IsSuccess(), createTableResult.GetIssues().ToString());
+
+            auto result = session.ExecuteDataQuery(Q1_(R"(
+                UPSERT INTO `/Root/Right` (Key, Value) VALUES (1, 'rValue1'), (3, 'rValue3'), (NULL, 'rValue');
+            )"), TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {  // inner
+            const auto query = Q1_(R"(
+                SELECT l.Value, r.Value FROM `/Root/Left` AS l JOIN `/Root/Right` AS r ON l.Key = r.Key;
+            )");
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            CompareYson(R"([[["lValue1"];["rValue1"]]])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+
+        {  // left
+            const auto query = Q1_(R"(
+                SELECT l.Value, r.Value FROM `/Root/Left` AS l LEFT JOIN `/Root/Right` AS r ON l.Key = r.Key ORDER BY l.Value;
+            )");
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            CompareYson(R"([[["lValue1"];["rValue1"]];[["lValue2"];#]])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+
+        {  // right
+            const auto query = Q1_(R"(
+                SELECT r.Value, l.Value FROM `/Root/Left` AS l RIGHT JOIN `/Root/Right` AS r ON l.Key = r.Key ORDER BY r.Value;
+            )");
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            CompareYson(R"([[["rValue"];#];[["rValue1"];["lValue1"]];[["rValue3"];#]])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+    }
 }
 
 } // namespace NKqp