Move YDB-specific code to YDB

5 files changed, 58 insertions, 52 deletions

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
index 118442443d..fe6c650591 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
@@ -173,7 +173,8 @@ TExprBase KqpRewriteSqlInToEquiJoin(const TExprBase& node, TExprContext& ctx, co
             return false;
         }
 
-        return CanRewriteSqlInToEquiJoin(sqlIn.Lookup().Ref().GetTypeAnn(), sqlIn.Collection().Ref().GetTypeAnn());
+        return CanRewriteSqlInToEquiJoin(sqlIn.Lookup().Ref().GetTypeAnn(), sqlIn.Collection().Ref().GetTypeAnn())
+            && CanConvertSqlInToJoin(sqlIn);
     };
 
     const bool prefixOnly = true;
diff --git a/ydb/core/kqp/provider/yql_kikimr_opt_range.cpp b/ydb/core/kqp/provider/yql_kikimr_opt_range.cpp
index cce81efc42..dcbf34afd8 100644
--- a/ydb/core/kqp/provider/yql_kikimr_opt_range.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_opt_range.cpp
@@ -1031,7 +1031,8 @@ TExprNode::TPtr KiSqlInToEquiJoin(NNodes::TExprBase node, const TKikimrTablesDat
             return false;
         }
 
-        return CanRewriteSqlInToEquiJoin(sqlIn.Lookup().Ref().GetTypeAnn(), sqlIn.Collection().Ref().GetTypeAnn());
+        return CanRewriteSqlInToEquiJoin(sqlIn.Lookup().Ref().GetTypeAnn(), sqlIn.Collection().Ref().GetTypeAnn())
+            && CanConvertSqlInToJoin(sqlIn);
     };
 
     const bool prefixOnly = true;
diff --git a/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp b/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
index 912dfca53c..12636d89d4 100644
--- a/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_opt_utils.cpp
@@ -69,4 +69,53 @@ bool CanRewriteSqlInToEquiJoin(const TTypeAnnotationNode* lookupType, const TTyp
     return false;
 }
 
+bool CanConvertSqlInToJoin(const NNodes::TCoSqlIn& sqlIn) {
+    auto leftArg = sqlIn.Lookup();
+    auto leftColumnType = leftArg.Ref().GetTypeAnn();
+
+    auto rightArg = sqlIn.Collection();
+    auto rightArgType = rightArg.Ref().GetTypeAnn();
+
+    if (rightArgType->GetKind() == ETypeAnnotationKind::List) {
+        auto rightListItemType = rightArgType->Cast<TListExprType>()->GetItemType();
+
+        auto isDataOrTupleOfData = [](const TTypeAnnotationNode* type) {
+            if (IsDataOrOptionalOfData(type) || type->GetKind() == ETypeAnnotationKind::Pg) {
+                return true;
+            }
+            if (type->GetKind() == ETypeAnnotationKind::Tuple) {
+                return AllOf(type->Cast<TTupleExprType>()->GetItems(), [](const auto& item) {
+                    return IsDataOrOptionalOfData(item);
+                });
+            }
+            return false;
+        };
+
+        if (rightListItemType->GetKind() == ETypeAnnotationKind::Struct) {
+            auto rightStructType = rightListItemType->Cast<TStructExprType>();
+            YQL_ENSURE(rightStructType->GetSize() == 1);
+            auto rightColumnType = rightStructType->GetItems()[0]->GetItemType();
+            return isDataOrTupleOfData(rightColumnType);
+        }
+
+        return isDataOrTupleOfData(rightListItemType);
+    }
+
+    /**
+     * todo: support tuple of equal tuples
+     *
+     * sql expression \code{.sql} ... where (k1, k2) in ((1, 2), (2, 3), (3, 4)) \endcode
+     * is equivalent to the \code{.sql} ... where (k1, k2) in AsTuple((1, 2), (2, 3), (3, 4)) \endcode
+     * but not to the \code{.sql} ... where (k1, k2) in AsList((1, 2), (2, 3), (3, 4)) \endcode
+     * so, it's not supported now
+     */
+
+    if (rightArgType->GetKind() == ETypeAnnotationKind::Dict) {
+        auto rightDictType = rightArgType->Cast<TDictExprType>()->GetKeyType();
+        return IsDataOrOptionalOfData(leftColumnType) && IsDataOrOptionalOfData(rightDictType);
+    }
+
+    return false;
+}
+
 } // namespace NYql
diff --git a/ydb/core/kqp/provider/yql_kikimr_opt_utils.h b/ydb/core/kqp/provider/yql_kikimr_opt_utils.h
index 219a36f5e8..01da33caff 100644
--- a/ydb/core/kqp/provider/yql_kikimr_opt_utils.h
+++ b/ydb/core/kqp/provider/yql_kikimr_opt_utils.h
@@ -2,6 +2,8 @@
 
 #include "yql_kikimr_provider.h"
 
+#include <ydb/library/yql/core/common_opt/yql_co_sqlin.h>
+
 namespace NYql {
 
 bool GetEquiJoinKeyTypes(NNodes::TExprBase leftInput, const TString& leftColumnName,
@@ -10,4 +12,6 @@ bool GetEquiJoinKeyTypes(NNodes::TExprBase leftInput, const TString& leftColumnN
 
 bool CanRewriteSqlInToEquiJoin(const TTypeAnnotationNode* lookupType, const TTypeAnnotationNode* collectionType);
 
+bool CanConvertSqlInToJoin(const NNodes::TCoSqlIn& sqlIn);
+
 } // namespace NYql
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index 484f8f3697..a0f5d4e15b 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -1461,55 +1461,6 @@ bool ShouldConvertSqlInToJoin(const TCoSqlIn& sqlIn, bool /* negated */) {
     return tableSource;
 }
 
-bool CanConvertSqlInToJoin(const TCoSqlIn& sqlIn) {
-    auto leftArg = sqlIn.Lookup();
-    auto leftColumnType = leftArg.Ref().GetTypeAnn();
-
-    auto rightArg = sqlIn.Collection();
-    auto rightArgType = rightArg.Ref().GetTypeAnn();
-
-    if (rightArgType->GetKind() == ETypeAnnotationKind::List) {
-        auto rightListItemType = rightArgType->Cast<TListExprType>()->GetItemType();
-
-        auto isDataOrTupleOfData = [](const TTypeAnnotationNode* type) {
-            if (IsDataOrOptionalOfData(type) || type->GetKind() == ETypeAnnotationKind::Pg) {
-                return true;
-            }
-            if (type->GetKind() == ETypeAnnotationKind::Tuple) {
-                return AllOf(type->Cast<TTupleExprType>()->GetItems(), [](const auto& item) {
-                    return IsDataOrOptionalOfData(item);
-                });
-            }
-            return false;
-        };
-
-        if (rightListItemType->GetKind() == ETypeAnnotationKind::Struct) {
-            auto rightStructType = rightListItemType->Cast<TStructExprType>();
-            YQL_ENSURE(rightStructType->GetSize() == 1);
-            auto rightColumnType = rightStructType->GetItems()[0]->GetItemType();
-            return isDataOrTupleOfData(rightColumnType);
-        }
-
-        return isDataOrTupleOfData(rightListItemType);
-    }
-
-    /**
-     * todo: support tuple of equal tuples
-     *
-     * sql expression \code{.sql} ... where (k1, k2) in ((1, 2), (2, 3), (3, 4)) \endcode
-     * is equivalent to the \code{.sql} ... where (k1, k2) in AsTuple((1, 2), (2, 3), (3, 4)) \endcode
-     * but not to the \code{.sql} ... where (k1, k2) in AsList((1, 2), (2, 3), (3, 4)) \endcode
-     * so, it's not supported now
-     */
-
-    if (rightArgType->GetKind() == ETypeAnnotationKind::Dict) {
-        auto rightDictType = rightArgType->Cast<TDictExprType>()->GetKeyType();
-        return IsDataOrOptionalOfData(leftColumnType) && IsDataOrOptionalOfData(rightDictType);
-    }
-
-    return false;
-}
-
 struct TPredicateChainNode {
     TExprNode::TPtr Pred;
 
@@ -1674,7 +1625,7 @@ TPredicateChainNode ParsePredicateChainNode(const TExprNode::TPtr& predicate, co
     }
 
     TCoSqlIn sqlIn(curr);
-    if (!shouldConvertSqlInToJoin(sqlIn, result.Negated) || !CanConvertSqlInToJoin(sqlIn)) {
+    if (!shouldConvertSqlInToJoin(sqlIn, result.Negated)) {
         // not convertible to join
         return result;
     }