[embedded] Full plan

YQL-15884

6 files changed, 33 insertions, 14 deletions

diff --git a/ydb/library/yql/core/facade/yql_facade.cpp b/ydb/library/yql/core/facade/yql_facade.cpp
index 7cfae4e7516..3e4d699722e 100644
--- a/ydb/library/yql/core/facade/yql_facade.cpp
+++ b/ydb/library/yql/core/facade/yql_facade.cpp
@@ -1064,7 +1064,7 @@ TMaybe<TString> TProgram::GetQueryAst() {
     return Nothing();
 }
 
-TMaybe<TString> TProgram::GetQueryPlan() {
+TMaybe<TString> TProgram::GetQueryPlan(const TPlanSettings& settings) {
     if (ExternalQueryPlan_) {
         return ExternalQueryPlan_;
     }
@@ -1074,7 +1074,7 @@ TMaybe<TString> TProgram::GetQueryPlan() {
         planStream.Reserve(DEFAULT_PLAN_BUF_SIZE);
 
         NYson::TYsonWriter writer(&planStream, OutputFormat_);
-        PlanBuilder_->WritePlan(writer, ExprRoot_);
+        PlanBuilder_->WritePlan(writer, ExprRoot_, settings);
 
         return planStream.Str();
     }
diff --git a/ydb/library/yql/core/facade/yql_facade.h b/ydb/library/yql/core/facade/yql_facade.h
index 08062c06235..8a62b6604f6 100644
--- a/ydb/library/yql/core/facade/yql_facade.h
+++ b/ydb/library/yql/core/facade/yql_facade.h
@@ -213,7 +213,7 @@ public:
     }
 
     TMaybe<TString> GetQueryAst();
-    TMaybe<TString> GetQueryPlan();
+    TMaybe<TString> GetQueryPlan(const TPlanSettings& settings = {});
 
     void SetDiagnosticFormat(NYson::EYsonFormat format) {
         DiagnosticFormat_ = format;
diff --git a/ydb/library/yql/core/services/yql_out_transformers.cpp b/ydb/library/yql/core/services/yql_out_transformers.cpp
index d8363095f89..5e4cfccb9ed 100644
--- a/ydb/library/yql/core/services/yql_out_transformers.cpp
+++ b/ydb/library/yql/core/services/yql_out_transformers.cpp
@@ -37,11 +37,11 @@ IGraphTransformer::TStatus TPlanOutputTransformer::operator()(
     output = input;
     if (DirectOut_) {
         NYson::TYsonWriter writer(DirectOut_, OutputFormat_);
-        Builder_.WritePlan(writer, input);
+        Builder_.WritePlan(writer, input, PlanSettings_);
     } else {
         TNullOutput null;
         NYson::TYsonWriter writer(&null, OutputFormat_);
-        Builder_.WritePlan(writer, input);
+        Builder_.WritePlan(writer, input, PlanSettings_);
     }
 
     return IGraphTransformer::TStatus::Ok;
diff --git a/ydb/library/yql/core/services/yql_out_transformers.h b/ydb/library/yql/core/services/yql_out_transformers.h
index d90f88cd997..b1e4c39b5a9 100644
--- a/ydb/library/yql/core/services/yql_out_transformers.h
+++ b/ydb/library/yql/core/services/yql_out_transformers.h
@@ -41,10 +41,12 @@ public:
     TPlanOutputTransformer(
         IOutputStream* directOut,
         IPlanBuilder& builder,
-        NYson::EYsonFormat outputFormat)
+        NYson::EYsonFormat outputFormat,
+        TPlanSettings&& settings = {})
         : DirectOut_(directOut)
         , Builder_(builder)
         , OutputFormat_(outputFormat)
+        , PlanSettings_(std::move(settings))
     {
     }
 
@@ -62,6 +64,7 @@ private:
     IOutputStream* DirectOut_;
     IPlanBuilder& Builder_;
     NYson::EYsonFormat OutputFormat_;
+    TPlanSettings PlanSettings_;
 };
 
 class TExprLogTransformer {
diff --git a/ydb/library/yql/core/services/yql_plan.cpp b/ydb/library/yql/core/services/yql_plan.cpp
index 4d7cfffaf5a..c29c88f588e 100644
--- a/ydb/library/yql/core/services/yql_plan.cpp
+++ b/ydb/library/yql/core/services/yql_plan.cpp
@@ -17,8 +17,6 @@ struct TPinAttrs {
     {}
 };
 
-const ui32 LIMIT_PINS = 10;
-
 struct TNodeInfo {
     ui64 NodeId;
     const TExprNode* const Node;
@@ -184,7 +182,7 @@ public:
         : Types_(types)
     {}
 
-    void WritePlan(NYson::TYsonWriter& writer, const TExprNode::TPtr& root) override {
+    void WritePlan(NYson::TYsonWriter& writer, const TExprNode::TPtr& root, const TPlanSettings& settings) override {
         if (!root) {
             return;
         }
@@ -234,12 +232,12 @@ public:
 
                 info.InputsCount = inputs.size();
                 info.OutputsCount = outputs.size();
-                if (inputs.size() > LIMIT_PINS) {
-                    inputs.resize(LIMIT_PINS, TPinInfo(nullptr, nullptr, nullptr, "", true));
+                if (settings.LimitInputPins && inputs.size() > *settings.LimitInputPins) {
+                    inputs.resize(*settings.LimitInputPins, TPinInfo(nullptr, nullptr, nullptr, "", true));
                 }
 
-                if (outputs.size() > LIMIT_PINS) {
-                    outputs.resize(LIMIT_PINS, TPinInfo(nullptr, nullptr, nullptr, "", true));
+                if (settings.LimitOutputPins && outputs.size() > *settings.LimitOutputPins) {
+                    outputs.resize(*settings.LimitOutputPins, TPinInfo(nullptr, nullptr, nullptr, "", true));
                 }
 
                 WritePins("Inputs", inputs, writer, info.Inputs, providers);
diff --git a/ydb/library/yql/core/services/yql_plan.h b/ydb/library/yql/core/services/yql_plan.h
index 608c54fb8e4..2c7f19fdd37 100644
--- a/ydb/library/yql/core/services/yql_plan.h
+++ b/ydb/library/yql/core/services/yql_plan.h
@@ -4,13 +4,31 @@
 
 #include <library/cpp/yson/writer.h>
 
+#include <optional>
+
 namespace NYql {
 
+struct TPlanSettings {
+
+    TPlanSettings& SetLimitInputPins(std::optional<ui32> val) {
+        LimitInputPins = std::move(val);
+        return *this;
+    }
+
+    TPlanSettings& SetLimitOutputPins(std::optional<ui32> val) {
+        LimitOutputPins = std::move(val);
+        return *this;
+    }
+
+    std::optional<ui32> LimitInputPins = 10;
+    std::optional<ui32> LimitOutputPins = 10;
+};
+
 class IPlanBuilder {
 public:
     virtual ~IPlanBuilder() {};
     virtual void Clear() = 0;
-    virtual void WritePlan(NYson::TYsonWriter& writer, const TExprNode::TPtr& root) = 0;
+    virtual void WritePlan(NYson::TYsonWriter& writer, const TExprNode::TPtr& root, const TPlanSettings& settings = {}) = 0;
 };
 
 TAutoPtr<IPlanBuilder> CreatePlanBuilder(TTypeAnnotationContext& types);