Fix YT provider direct dependency on DQ via CBO

nepal · nepal · commit 340bba021b6b · 2024-11-26T18:00:13.000+03:00
commit_hash:134201d428f9541bb6990e2d65a7bbcf65d74e25
diff --git a/yql/essentials/core/cbo/cbo_optimizer_new.h b/yql/essentials/core/cbo/cbo_optimizer_new.h
@@ -76,7 +76,7 @@ struct TCardinalityHints {
 
         double ApplyHint(double originalValue) {
             Applied = true;
-            
+
             switch (Operation) {
                 case Add:
                     return originalValue + Value;
@@ -122,7 +122,7 @@ struct TJoinOrderHints {
 
         virtual TVector<TString> Labels() = 0;
         bool IsRelation() { return Type == Relation; }
-        bool IsJoin() { return Type == Join; } 
+        bool IsJoin() { return Type == Join; }
         virtual ~ITreeNode() = default;
 
         ui32 Type;
@@ -136,7 +136,7 @@ struct TJoinOrderHints {
             this->Type = ITreeNode::Join;
         }
 
-        TVector<TString> Labels() override {     
+        TVector<TString> Labels() override {
             auto labels = Lhs->Labels();
             auto rhsLabels = Rhs->Labels();
             labels.insert(labels.end(), std::make_move_iterator(rhsLabels.begin()), std::make_move_iterator(rhsLabels.end()));
@@ -179,12 +179,12 @@ struct TOptimizerHints {
 
     TVector<TString> GetUnappliedString();
 
-    /* 
+    /*
      *   The function accepts string with three type of expressions: array of (JoinAlgo | Card | JoinOrder):
      *   1) JoinAlgo(t1 t2 ... tn Map | Grace | Lookup) to change join algo for join, where these labels take part
      *   2) Card(t1 t2 ... tn (*|/|+|-) Number) to change cardinality for join, where these labels take part or labels only
      *   3) JoinOrder( (t1 t2) (t3 (t4 ...)) ) - fixate this join subtree in the general join tree
-     */  
+     */
     static TOptimizerHints Parse(const TString&);
 };
 
@@ -299,14 +299,35 @@ struct TJoinOptimizerNode : public IBaseOptimizerNode {
 };
 
 struct IOptimizerNew {
+    using TPtr = std::shared_ptr<IOptimizerNew>;
     IProviderContext& Pctx;
 
     IOptimizerNew(IProviderContext& ctx) : Pctx(ctx) {}
     virtual ~IOptimizerNew() = default;
     virtual std::shared_ptr<TJoinOptimizerNode> JoinSearch(
-        const std::shared_ptr<TJoinOptimizerNode>& joinTree, 
+        const std::shared_ptr<TJoinOptimizerNode>& joinTree,
         const TOptimizerHints& hints = {}
     ) = 0;
 };
 
+struct TExprContext;
+
+class IOptimizerFactory : private TNonCopyable {
+public:
+    using TPtr = std::shared_ptr<IOptimizerFactory>;
+    using TLogger = std::function<void(const TString&)>;
+
+    virtual ~IOptimizerFactory() = default;
+
+    struct TNativeSettings {
+        ui32 MaxDPhypDPTableSize = 100000;
+    };
+    virtual IOptimizerNew::TPtr MakeJoinCostBasedOptimizerNative(IProviderContext& pctx, TExprContext& ctx, const TNativeSettings& settings) const = 0;
+
+    struct TPGSettings {
+        TLogger Logger = [](const TString&){};
+    };
+    virtual IOptimizerNew::TPtr MakeJoinCostBasedOptimizerPG(IProviderContext& pctx, TExprContext& ctx, const TPGSettings& settings) const = 0;
+};
+
 } // namespace NYql
diff --git a/yql/essentials/core/cbo/simple/cbo_simple.cpp b/yql/essentials/core/cbo/simple/cbo_simple.cpp
@@ -0,0 +1,31 @@
+#include "cbo_simple.h"
+
+#include <yql/essentials/utils/yql_panic.h>
+#include <yql/essentials/parser/pg_wrapper/interface/optimizer.h>
+
+namespace NYql {
+
+namespace {
+
+class TSimpleOptimizerFactory : public IOptimizerFactory {
+public:
+    virtual IOptimizerNew::TPtr MakeJoinCostBasedOptimizerNative(IProviderContext& pctx, TExprContext& ctx, const TNativeSettings& settings) const override {
+        Y_UNUSED(pctx);
+        Y_UNUSED(ctx);
+        Y_UNUSED(settings);
+        YQL_ENSURE(false, "Native CBO is not supported here");
+        Y_UNREACHABLE();
+    }
+
+    virtual IOptimizerNew::TPtr MakeJoinCostBasedOptimizerPG(IProviderContext& pctx, TExprContext& ctx, const TPGSettings& settings) const override {
+        return IOptimizerNew::TPtr(MakePgOptimizerNew(pctx, ctx, settings.Logger));
+    }
+};
+
+}
+
+IOptimizerFactory::TPtr MakeSimpleCBOOptimizerFactory() {
+    return std::make_shared<TSimpleOptimizerFactory>();
+}
+
+} // namespace NYql
diff --git a/yql/essentials/core/cbo/simple/cbo_simple.h b/yql/essentials/core/cbo/simple/cbo_simple.h
@@ -0,0 +1,9 @@
+#pragma once
+
+#include <yql/essentials/core/cbo/cbo_optimizer_new.h>
+
+namespace NYql {
+
+IOptimizerFactory::TPtr MakeSimpleCBOOptimizerFactory();
+
+} // namespace NYql
diff --git a/yql/essentials/core/cbo/simple/ya.make b/yql/essentials/core/cbo/simple/ya.make
@@ -0,0 +1,12 @@
+LIBRARY()
+
+SRCS(
+    cbo_simple.cpp
+)
+
+PEERDIR(
+    yql/essentials/parser/pg_wrapper/interface
+    yql/essentials/utils
+)
+
+END()
diff --git a/yql/essentials/core/cbo/ya.make b/yql/essentials/core/cbo/ya.make
@@ -9,7 +9,10 @@ GENERATE_ENUM_SERIALIZATION(cbo_optimizer_new.h)
 
 END()
 
+RECURSE(
+    simple
+)
+
 RECURSE_FOR_TESTS(
     ut
 )
-
diff --git a/yql/essentials/core/extract_predicate/ut/extract_predicate_ut.cpp b/yql/essentials/core/extract_predicate/ut/extract_predicate_ut.cpp
@@ -10,6 +10,7 @@
 #include <yql/essentials/sql/sql.h>
 #include <yql/essentials/ast/yql_ast_annotation.h>
 #include <yql/essentials/ast/yql_expr.h>
+#include <yql/essentials/core/cbo/simple/cbo_simple.h>
 #include <yql/essentials/core/type_ann/type_ann_core.h>
 #include <yql/essentials/core/type_ann/type_ann_expr.h>
 #include <yql/essentials/core/ut_common/yql_ut_common.h>
@@ -154,7 +155,7 @@ Y_UNIT_TEST_SUITE(TYqlExtractPredicate) {
             auto ytGateway = CreateYtFileGateway(yqlNativeServices);
 
             TVector<TDataProviderInitializer> dataProvidersInit;
-            dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway));
+            dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway, MakeSimpleCBOOptimizerFactory()));
             TProgramFactory factory(true, funcReg, 0ULL, dataProvidersInit, "ut");
 
             TProgramPtr program = factory.Create("-stdin-", Src);
diff --git a/yql/essentials/core/extract_predicate/ut/ya.make b/yql/essentials/core/extract_predicate/ut/ya.make
@@ -22,6 +22,7 @@ PEERDIR(
     contrib/ydb/library/yql/providers/yt/provider
     contrib/ydb/library/yql/providers/yt/codec/codegen
     contrib/ydb/library/yql/providers/yt/comp_nodes/llvm14
+    yql/essentials/core/cbo/simple
     yql/essentials/minikql/comp_nodes/llvm14
     yql/essentials/minikql/invoke_builtins/llvm14
     yql/essentials/sql/pg
diff --git a/yql/essentials/core/ut/ya.make b/yql/essentials/core/ut/ya.make
@@ -22,6 +22,7 @@ PEERDIR(
     library/cpp/yson/node
     yql/essentials/ast
     yql/essentials/core
+    yql/essentials/core/cbo/simple
     yql/essentials/core/facade
     yql/essentials/core/services
     yql/essentials/core/services/mounts
diff --git a/yql/essentials/core/ut/yql_execution_ut.cpp b/yql/essentials/core/ut/yql_execution_ut.cpp
@@ -4,6 +4,7 @@
 
 #include <yql/essentials/ast/yql_ast_annotation.h>
 #include <yql/essentials/ast/yql_expr.h>
+#include <yql/essentials/core/cbo/simple/cbo_simple.h>
 #include <yql/essentials/core/type_ann/type_ann_core.h>
 #include <yql/essentials/core/yql_expr_optimize.h>
 #include <yql/essentials/core/yql_expr_type_annotation.h>
@@ -58,7 +59,7 @@ namespace NYql {
             auto ytGateway = CreateYtFileGateway(yqlNativeServices);
 
             TVector<TDataProviderInitializer> dataProvidersInit;
-            dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway));
+            dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway, MakeSimpleCBOOptimizerFactory()));
             TProgramFactory factory(true, funcReg, 0ULL, dataProvidersInit, "ut");
 
             TProgramPtr program = factory.Create("-stdin-", Src);
diff --git a/yql/essentials/core/ut/yql_qplayer_ut.cpp b/yql/essentials/core/ut/yql_qplayer_ut.cpp
@@ -2,6 +2,7 @@
 #include <contrib/ydb/library/yql/providers/yt/provider/yql_yt_provider.h>
 #include <contrib/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.h>
 #include <contrib/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.h>
+#include <yql/essentials/core/cbo/simple/cbo_simple.h>
 #include <yql/essentials/core/facade/yql_facade.h>
 #include <yql/essentials/core/qplayer/storage/memory/yql_qstorage_memory.h>
 #include <yql/essentials/providers/common/udf_resolve/yql_simple_udf_resolver.h>
@@ -82,7 +83,7 @@ bool RunProgram(bool replay, const TString& query, const TQContext& qContext, co
     auto ytGateway = CreateYtFileGateway(yqlNativeServices);
 
     TVector<TDataProviderInitializer> dataProvidersInit;
-    dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway));
+    dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytGateway, MakeSimpleCBOOptimizerFactory()));
 
     TExprContext modulesCtx;
     IModuleResolver::TPtr moduleResolver;
@@ -113,7 +114,7 @@ bool RunProgram(bool replay, const TString& query, const TQContext& qContext, co
         for (const auto& x : runSettings.Credentials) {
             credentials->AddCredential(x.first, x.second);
         }
-        
+
         factory.SetCredentials(credentials);
     }
 
diff --git a/yql/essentials/tools/pgrun/pgrun.cpp b/yql/essentials/tools/pgrun/pgrun.cpp
@@ -1,6 +1,7 @@
 #include <yql/essentials/utils/backtrace/backtrace.h>
 #include <yql/essentials/minikql/invoke_builtins/mkql_builtins.h>
 #include <yql/essentials/minikql/mkql_function_registry.h>
+#include <yql/essentials/core/cbo/simple/cbo_simple.h>
 #include <yql/essentials/core/facade/yql_facade.h>
 #include <yql/essentials/core/yql_opt_utils.h>
 #include <yql/essentials/core/yql_expr_optimize.h>
@@ -1169,7 +1170,7 @@ int Main(int argc, char* argv[])
     }
 
     TVector<TDataProviderInitializer> dataProvidersInit;
-    dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytNativeGateway));
+    dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytNativeGateway, MakeSimpleCBOOptimizerFactory()));
     dataProvidersInit.push_back(GetPgDataProviderInitializer());
 
     TExprContext ctx;
diff --git a/yql/essentials/tools/pgrun/ya.make b/yql/essentials/tools/pgrun/ya.make
@@ -16,6 +16,7 @@ PEERDIR(
     library/cpp/yson
     library/cpp/yson/node
     yql/essentials/sql/pg
+    yql/essentials/core/cbo/simple
     yql/essentials/core/facade
     yql/essentials/core/file_storage
     yql/essentials/core/file_storage/proto
