Optimized DPccp (#1563)

* Optimized DPccp

* Cleaned up new optimizer interface

* Simplified the optimization a lot

* Added some comments

8 files changed, 290 insertions, 80 deletions

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp b/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp
index fed03ea383d..ef0604260eb 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp
@@ -90,9 +90,12 @@ bool IsLookupJoinApplicableDetailed(const std::shared_ptr<NYql::TRelOptimizerNod
 bool IsLookupJoinApplicable(std::shared_ptr<IBaseOptimizerNode> left, 
     std::shared_ptr<IBaseOptimizerNode> right, 
     const std::set<std::pair<TJoinColumn, TJoinColumn>>& joinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
     TKqpProviderContext& ctx) {
 
     Y_UNUSED(left);
+    Y_UNUSED(leftJoinKeys);
 
     auto rightStats = right->Stats;
 
@@ -114,12 +117,7 @@ bool IsLookupJoinApplicable(std::shared_ptr<IBaseOptimizerNode> left,
         }
     }
 
-    TVector<TString> joinKeys;
-    for( auto [leftJc, rightJc] : joinConditions ) {
-        joinKeys.emplace_back( rightJc.AttributeName);
-    }
-
-    return IsLookupJoinApplicableDetailed(std::static_pointer_cast<TRelOptimizerNode>(right), joinKeys, ctx);
+    return IsLookupJoinApplicableDetailed(std::static_pointer_cast<TRelOptimizerNode>(right), rightJoinKeys, ctx);
 }
 
 }
@@ -127,6 +125,8 @@ bool IsLookupJoinApplicable(std::shared_ptr<IBaseOptimizerNode> left,
 bool TKqpProviderContext::IsJoinApplicable(const std::shared_ptr<IBaseOptimizerNode>& left, 
     const std::shared_ptr<IBaseOptimizerNode>& right, 
     const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
     EJoinAlgoType joinAlgo)  {
 
     switch( joinAlgo ) {
@@ -134,7 +134,7 @@ bool TKqpProviderContext::IsJoinApplicable(const std::shared_ptr<IBaseOptimizerN
             if (OptLevel==2 && left->Stats->Nrows > 10e3) {
                 return false;
             }
-            return IsLookupJoinApplicable(left, right, joinConditions, *this);
+            return IsLookupJoinApplicable(left, right, joinConditions, leftJoinKeys, rightJoinKeys, *this);
 
         case EJoinAlgoType::DictJoin:
             return right->Stats->Nrows < 10e5;
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_cbo.h b/ydb/core/kqp/opt/logical/kqp_opt_cbo.h
index 13b6b0200ec..e43e0a33c19 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_cbo.h
+++ b/ydb/core/kqp/opt/logical/kqp_opt_cbo.h
@@ -26,6 +26,7 @@ struct TKqpProviderContext : public NYql::IProviderContext {
     virtual bool IsJoinApplicable(const std::shared_ptr<NYql::IBaseOptimizerNode>& left, 
         const std::shared_ptr<NYql::IBaseOptimizerNode>& right, 
         const std::set<std::pair<NYql::NDq::TJoinColumn, NYql::NDq::TJoinColumn>>& joinConditions,
+        const TVector<TString>& leftJoinKeys, const TVector<TString>& rightJoinKeys,
         NYql::EJoinAlgoType joinAlgo) override;
 
     virtual double ComputeJoinCost(const NYql::TOptimizerStatistics& leftStats, const NYql::TOptimizerStatistics& rightStats, NYql::EJoinAlgoType joinAlgo) const override;
diff --git a/ydb/library/yql/core/cbo/cbo_optimizer_new.cpp b/ydb/library/yql/core/cbo/cbo_optimizer_new.cpp
index c7ce233d76f..c125cb5bc00 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer_new.cpp
+++ b/ydb/library/yql/core/cbo/cbo_optimizer_new.cpp
@@ -67,10 +67,14 @@ TJoinOptimizerNode::TJoinOptimizerNode(const std::shared_ptr<IBaseOptimizerNode>
     IBaseOptimizerNode(JoinNodeType), 
     LeftArg(left), 
     RightArg(right), 
-    JoinConditions(joinConditions), 
+    JoinConditions(joinConditions),
     JoinType(joinType),
     JoinAlgo(joinAlgo) {
         IsReorderable = (JoinType==EJoinKind::InnerJoin) && (nonReorderable==false);
+        for (auto [l,r] : joinConditions ) {
+            LeftJoinKeys.push_back(l.AttributeName);
+            RightJoinKeys.push_back(r.AttributeName);
+        }
     }
 
 TVector<TString> TJoinOptimizerNode::Labels() {
@@ -97,7 +101,9 @@ void TJoinOptimizerNode::Print(std::stringstream& stream, int ntabs) {
         stream << "\t";
     }
 
-    stream << *Stats << "\n";
+    if (Stats) {
+        stream << *Stats << "\n";
+    }
 
     LeftArg->Print(stream, ntabs+1);
     RightArg->Print(stream, ntabs+1);
diff --git a/ydb/library/yql/core/cbo/cbo_optimizer_new.h b/ydb/library/yql/core/cbo/cbo_optimizer_new.h
index 256252241e7..e4c53da4fce 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer_new.h
+++ b/ydb/library/yql/core/cbo/cbo_optimizer_new.h
@@ -92,6 +92,8 @@ struct IProviderContext {
     virtual bool IsJoinApplicable(const std::shared_ptr<IBaseOptimizerNode>& left, 
         const std::shared_ptr<IBaseOptimizerNode>& right, 
         const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+        const TVector<TString>& leftJoinKeys,
+        const TVector<TString>& rightJoinKeys,
         EJoinAlgoType joinAlgo) = 0;
 
 };
@@ -111,11 +113,15 @@ struct TDummyProviderContext : public IProviderContext {
     bool IsJoinApplicable(const std::shared_ptr<IBaseOptimizerNode>& left, 
         const std::shared_ptr<IBaseOptimizerNode>& right, 
         const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+        const TVector<TString>& leftJoinKeys,
+        const TVector<TString>& rightJoinKeys,
         EJoinAlgoType joinAlgo) override {
 
         Y_UNUSED(left);
         Y_UNUSED(right);
         Y_UNUSED(joinConditions);
+        Y_UNUSED(leftJoinKeys);
+        Y_UNUSED(rightJoinKeys);
         Y_UNUSED(joinAlgo);
 
         return true;
@@ -137,13 +143,19 @@ struct TDummyProviderContext : public IProviderContext {
 struct TJoinOptimizerNode : public IBaseOptimizerNode {
     std::shared_ptr<IBaseOptimizerNode> LeftArg;
     std::shared_ptr<IBaseOptimizerNode> RightArg;
-    std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>> JoinConditions;
+    const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>> JoinConditions;
+    TVector<TString> LeftJoinKeys;
+    TVector<TString> RightJoinKeys;
     EJoinKind JoinType;
     EJoinAlgoType JoinAlgo;
     bool IsReorderable;
 
-    TJoinOptimizerNode(const std::shared_ptr<IBaseOptimizerNode>& left, const std::shared_ptr<IBaseOptimizerNode>& right, 
-        const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions, const EJoinKind joinType, const EJoinAlgoType joinAlgo, bool nonReorderable=false);
+    TJoinOptimizerNode(const std::shared_ptr<IBaseOptimizerNode>& left, 
+        const std::shared_ptr<IBaseOptimizerNode>& right, 
+        const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+        const EJoinKind joinType, 
+        const EJoinAlgoType joinAlgo, 
+        bool nonReorderable=false);
     virtual ~TJoinOptimizerNode() {}
     virtual TVector<TString> Labels();
     virtual void Print(std::stringstream& stream, int ntabs=0);
diff --git a/ydb/library/yql/core/yql_cost_function.cpp b/ydb/library/yql/core/yql_cost_function.cpp
index dcf395ca408..aae915e85d0 100644
--- a/ydb/library/yql/core/yql_cost_function.cpp
+++ b/ydb/library/yql/core/yql_cost_function.cpp
@@ -43,11 +43,13 @@ TOptimizerStatistics NYql::ComputeJoinStats(const TOptimizerStatistics& leftStat
 
     double newCard;
     EStatisticsType outputType;
-    TVector<TString> joinedTableKeys;
+    bool leftKeyColumns = false;
+    bool rightKeyColumns = false;
+
 
     if (IsPKJoin(rightStats,rightJoinKeys)) {
         newCard = leftStats.Nrows;
-        joinedTableKeys = leftStats.KeyColumns;
+        leftKeyColumns = true;
         if (leftStats.Type == EStatisticsType::BaseTable){
             outputType = EStatisticsType::FilteredFactTable;
         } else {
@@ -56,7 +58,7 @@ TOptimizerStatistics NYql::ComputeJoinStats(const TOptimizerStatistics& leftStat
     }
     else if (IsPKJoin(leftStats,leftJoinKeys)) {
         newCard = rightStats.Nrows;
-        joinedTableKeys = rightStats.KeyColumns;
+        rightKeyColumns = true;
         if (rightStats.Type == EStatisticsType::BaseTable){
             outputType = EStatisticsType::FilteredFactTable;
         } else {
@@ -74,9 +76,11 @@ TOptimizerStatistics NYql::ComputeJoinStats(const TOptimizerStatistics& leftStat
         + newCard 
         + leftStats.Cost + rightStats.Cost;
 
-    return TOptimizerStatistics(outputType, newCard, newNCols, cost, joinedTableKeys);
+    return TOptimizerStatistics(outputType, newCard, newNCols, cost, 
+        leftKeyColumns ? leftStats.KeyColumns : ( rightKeyColumns ? rightStats.KeyColumns : TOptimizerStatistics::EmptyColumns));
 }
 
+
 TOptimizerStatistics NYql::ComputeJoinStats(const TOptimizerStatistics& leftStats, const TOptimizerStatistics& rightStats, 
     const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions, EJoinAlgoType joinAlgo, const IProviderContext& ctx) {
 
@@ -90,3 +94,4 @@ TOptimizerStatistics NYql::ComputeJoinStats(const TOptimizerStatistics& leftStat
 
     return ComputeJoinStats(leftStats, rightStats, leftJoinKeys, rightJoinKeys, joinAlgo, ctx);
 }
+
diff --git a/ydb/library/yql/core/yql_statistics.cpp b/ydb/library/yql/core/yql_statistics.cpp
index 24033b50880..18f4d463dcd 100644
--- a/ydb/library/yql/core/yql_statistics.cpp
+++ b/ydb/library/yql/core/yql_statistics.cpp
@@ -17,3 +17,5 @@ TOptimizerStatistics& TOptimizerStatistics::operator+=(const TOptimizerStatistic
     Cost += other.Cost;
     return *this;
 }
+
+const TVector<TString>& TOptimizerStatistics::EmptyColumns = TVector<TString>();
diff --git a/ydb/library/yql/core/yql_statistics.h b/ydb/library/yql/core/yql_statistics.h
index 77eff37837d..a8969c18268 100644
--- a/ydb/library/yql/core/yql_statistics.h
+++ b/ydb/library/yql/core/yql_statistics.h
@@ -25,21 +25,19 @@ struct TOptimizerStatistics {
     double Nrows = 0;
     int Ncols = 0;
     double Cost;
-    TVector<TString> KeyColumns;
-
-    TString Descr;
-
-    TOptimizerStatistics() {}
-    TOptimizerStatistics(double nrows, int ncols): Nrows(nrows), Ncols(ncols) {}
-    TOptimizerStatistics(double nrows, int ncols, double cost): Nrows(nrows), Ncols(ncols), Cost(cost) {}
-    TOptimizerStatistics(EStatisticsType type, double nrows, int ncols, double cost): Type(type), Nrows(nrows), Ncols(ncols), Cost(cost) {}
-    TOptimizerStatistics(EStatisticsType type, double nrows, int ncols, double cost, TVector<TString> keyColumns): Type(type), Nrows(nrows), Ncols(ncols), Cost(cost), KeyColumns(keyColumns) {}
-    TOptimizerStatistics(double nrows,int ncols, double cost, TString descr): Nrows(nrows), Ncols(ncols), Cost(cost), Descr(descr) {}
+    const TVector<TString>& KeyColumns;
 
+    TOptimizerStatistics() : KeyColumns(EmptyColumns) {}
+    TOptimizerStatistics(double nrows, int ncols): Nrows(nrows), Ncols(ncols), KeyColumns(EmptyColumns) {}
+    TOptimizerStatistics(double nrows, int ncols, double cost): Nrows(nrows), Ncols(ncols), Cost(cost), KeyColumns(EmptyColumns) {}
+    TOptimizerStatistics(EStatisticsType type, double nrows, int ncols, double cost): Type(type), Nrows(nrows), Ncols(ncols), Cost(cost), KeyColumns(EmptyColumns) {}
+    TOptimizerStatistics(EStatisticsType type, double nrows, int ncols, double cost, const TVector<TString>& keyColumns): Type(type), Nrows(nrows), Ncols(ncols), Cost(cost), KeyColumns(keyColumns) {}
 
     TOptimizerStatistics& operator+=(const TOptimizerStatistics& other);
     bool Empty() const;
 
     friend std::ostream& operator<<(std::ostream& os, const TOptimizerStatistics& s);
+
+    static const TVector<TString>& EmptyColumns;
 };
 }
diff --git a/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp b/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
index 0ecc17f1532..48c778c76e4 100644
--- a/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
@@ -40,16 +40,30 @@ struct TEdge {
     int From;
     int To;
     mutable std::set<std::pair<TJoinColumn, TJoinColumn>> JoinConditions;
+    mutable TVector<TString> LeftJoinKeys;
+    mutable TVector<TString> RightJoinKeys;
+
+    TEdge(): From(-1), To(-1) {}
 
     TEdge(int f, int t): From(f), To(t) {}
 
     TEdge(int f, int t, std::pair<TJoinColumn, TJoinColumn> cond): From(f), To(t) {
         JoinConditions.insert(cond);
+        BuildCondVectors();
     }
 
     TEdge(int f, int t, std::set<std::pair<TJoinColumn, TJoinColumn>> conds): From(f), To(t), 
-        JoinConditions(conds) {}
+        JoinConditions(conds) {
+        BuildCondVectors();
+        }
 
+    void BuildCondVectors() {
+        for (auto [left, right] : JoinConditions) {
+            LeftJoinKeys.emplace_back(left.AttributeName);
+            RightJoinKeys.emplace_back(right.AttributeName);
+        }
+    }
+    
     bool operator==(const TEdge& other) const
     {
         return From==other.From && To==other.To;
@@ -62,6 +76,8 @@ struct TEdge {
             return e.From + e.To;
         }
     };
+
+    static const struct TEdge ErrorEdge;
 };
 
 /**
@@ -92,17 +108,148 @@ void ComputeJoinConditions(const TCoEquiJoinTuple& joinTuple,
 }
 
 /**
- * Create a new join and compute its statistics and cost
+ * Internal Join nodes are used inside the CBO. They don't own join condition data structures
+ * and therefore avoid copying them during generation of candidate plans.
+ * 
+ * These datastructures are owned by the query graph, so it is important to keep the graph around
+ * while internal nodes are being used.
+ * 
+ * After join enumeration, internal nodes need to be converted to regular nodes, that own the data
+ * structures
+*/
+struct TJoinOptimizerNodeInternal : public IBaseOptimizerNode {
+    std::shared_ptr<IBaseOptimizerNode> LeftArg;
+    std::shared_ptr<IBaseOptimizerNode> RightArg;
+    const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& JoinConditions;
+    const TVector<TString>& LeftJoinKeys;
+    const TVector<TString>& RightJoinKeys;
+    EJoinKind JoinType;
+    EJoinAlgoType JoinAlgo;
+    bool IsReorderable;
+
+    TJoinOptimizerNodeInternal(const std::shared_ptr<IBaseOptimizerNode>& left, 
+        const std::shared_ptr<IBaseOptimizerNode>& right, 
+        const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+        const TVector<TString>& leftJoinKeys,
+        const TVector<TString>& rightJoinKeys,
+        const EJoinKind joinType, 
+        const EJoinAlgoType joinAlgo, 
+        bool nonReorderable=false);
+
+    virtual ~TJoinOptimizerNodeInternal() {}
+    virtual TVector<TString> Labels();
+    virtual void Print(std::stringstream& stream, int ntabs=0);
+};
+
+TJoinOptimizerNodeInternal::TJoinOptimizerNodeInternal(const std::shared_ptr<IBaseOptimizerNode>& left, const std::shared_ptr<IBaseOptimizerNode>& right, 
+        const std::set<std::pair<TJoinColumn, TJoinColumn>>& joinConditions, const TVector<TString>& leftJoinKeys,
+        const TVector<TString>& rightJoinKeys, const EJoinKind joinType, const EJoinAlgoType joinAlgo, bool nonReorderable) :
+    IBaseOptimizerNode(JoinNodeType), 
+    LeftArg(left), 
+    RightArg(right), 
+    JoinConditions(joinConditions),
+    LeftJoinKeys(leftJoinKeys),
+    RightJoinKeys(rightJoinKeys),
+    JoinType(joinType),
+    JoinAlgo(joinAlgo) {
+        IsReorderable = (JoinType==EJoinKind::InnerJoin) && (nonReorderable==false);
+    }
+
+TVector<TString> TJoinOptimizerNodeInternal::Labels() {
+    auto res = LeftArg->Labels();
+    auto rightLabels = RightArg->Labels();
+    res.insert(res.begin(),rightLabels.begin(),rightLabels.end());
+    return res;
+}
+
+/**
+ * Convert a tree of internal optimizer nodes to external nodes that own the data structures.
+ * 
+ * The internal node tree can have references to external nodes (since some subtrees are optimized
+ * separately if the plan contains non-orderable joins). So we check the instances and if we encounter
+ * an external node, we return the whole subtree unchanged.
+*/
+std::shared_ptr<TJoinOptimizerNode> ConvertFromInternal(const std::shared_ptr<IBaseOptimizerNode> internal) {
+    Y_ENSURE(internal->Kind == EOptimizerNodeKind::JoinNodeType);
+    
+    if (dynamic_cast<TJoinOptimizerNode*>(internal.get()) != nullptr) {
+        return  std::static_pointer_cast<TJoinOptimizerNode>(internal);
+    }
+
+    auto join = std::static_pointer_cast<TJoinOptimizerNodeInternal>(internal);
+    
+    auto left = join->LeftArg;
+    auto right = join->RightArg;
+
+    if (left->Kind == EOptimizerNodeKind::JoinNodeType) {
+        left = ConvertFromInternal(left);
+    }
+    if (right->Kind == EOptimizerNodeKind::JoinNodeType) {
+        right = ConvertFromInternal(right);
+    }
+
+    auto newJoin = std::make_shared<TJoinOptimizerNode>(left, right, join->JoinConditions, join->JoinType, join->JoinAlgo, !join->IsReorderable);
+    newJoin->Stats = join->Stats;
+    return newJoin;
+}
+
+void TJoinOptimizerNodeInternal::Print(std::stringstream& stream, int ntabs) {
+    for (int i = 0; i < ntabs; i++){
+        stream << "\t";
+    }
+
+    stream << "Join: (" << JoinType << ") ";
+    for (auto c : JoinConditions){
+        stream << c.first.RelName << "." << c.first.AttributeName 
+            << "=" << c.second.RelName << "." 
+            << c.second.AttributeName << ", ";
+    }
+    stream << "\n";
+
+    for (int i = 0; i < ntabs; i++){
+        stream << "\t";
+    }
+
+    if (Stats) {
+        stream << *Stats << "\n";
+    }
+
+    LeftArg->Print(stream, ntabs+1);
+    RightArg->Print(stream, ntabs+1);
+}
+
+/**
+ * Create a new external join node and compute its statistics and cost
 */
 std::shared_ptr<TJoinOptimizerNode> MakeJoin(std::shared_ptr<IBaseOptimizerNode> left, 
     std::shared_ptr<IBaseOptimizerNode> right, 
     const std::set<std::pair<TJoinColumn, TJoinColumn>>& joinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
     EJoinKind joinKind,
     EJoinAlgoType joinAlgo,
+    bool nonReorderable,
     IProviderContext& ctx) {
 
-    auto res = std::make_shared<TJoinOptimizerNode>(left, right, joinConditions, joinKind, joinAlgo);
-    res->Stats = std::make_shared<TOptimizerStatistics>( ComputeJoinStats(*left->Stats, *right->Stats, joinConditions, joinAlgo, ctx));
+    auto res = std::make_shared<TJoinOptimizerNode>(left, right, joinConditions, joinKind, joinAlgo, nonReorderable);
+    res->Stats = std::make_shared<TOptimizerStatistics>( ComputeJoinStats(*left->Stats, *right->Stats, leftJoinKeys, rightJoinKeys, joinAlgo, ctx));
+    return res;
+}
+
+/**
+ * Create a new internal join node and compute its statistics and cost
+*/
+std::shared_ptr<TJoinOptimizerNodeInternal> MakeJoinInternal(std::shared_ptr<IBaseOptimizerNode> left, 
+    std::shared_ptr<IBaseOptimizerNode> right, 
+    const std::set<std::pair<TJoinColumn, TJoinColumn>>& joinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
+    EJoinKind joinKind,
+    EJoinAlgoType joinAlgo,
+    IProviderContext& ctx) {
+
+    auto res = std::make_shared<TJoinOptimizerNodeInternal>(left, right, joinConditions, leftJoinKeys, rightJoinKeys, joinKind, joinAlgo);
+    res->Stats = std::make_shared<TOptimizerStatistics>( ComputeJoinStats(*left->Stats, *right->Stats, leftJoinKeys, rightJoinKeys, joinAlgo, ctx));
     return res;
 }
 
@@ -110,44 +257,60 @@ std::shared_ptr<TJoinOptimizerNode> MakeJoin(std::shared_ptr<IBaseOptimizerNode>
  * Iterate over all join algorithms and pick the best join that is applicable.
  * Also considers commuting joins
 */
-std::shared_ptr<TJoinOptimizerNode> PickBestJoin(std::shared_ptr<IBaseOptimizerNode> left, 
+std::shared_ptr<TJoinOptimizerNodeInternal> PickBestJoin(std::shared_ptr<IBaseOptimizerNode> left, 
     std::shared_ptr<IBaseOptimizerNode> right, 
     const std::set<std::pair<TJoinColumn, TJoinColumn>>& leftJoinConditions,
     const std::set<std::pair<TJoinColumn, TJoinColumn>>& rightJoinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
     IProviderContext& ctx) {
 
-    auto res = std::shared_ptr<TJoinOptimizerNode>();
+    double bestCost;
+    EJoinAlgoType bestAlgo;
+    bool bestJoinLeftRight = true;
+    bool bestJoinValid = false;
 
     for ( auto joinAlgo : AllJoinAlgos ) {
-        auto p1 = ctx.IsJoinApplicable(left, right, leftJoinConditions, joinAlgo) ? 
-            MakeJoin(left, right, leftJoinConditions, EJoinKind::InnerJoin, joinAlgo, ctx) :
-            std::shared_ptr<TJoinOptimizerNode>();
-        auto p2 = ctx.IsJoinApplicable(right, left, rightJoinConditions, joinAlgo) ? 
-            MakeJoin(right, left, rightJoinConditions, EJoinKind::InnerJoin, joinAlgo, ctx) :
-            std::shared_ptr<TJoinOptimizerNode>();
-            
-        if (p1) {
-            if (res) {
-                if (p1->Stats->Cost < res->Stats->Cost) {
-                    res = p1;
+        if (ctx.IsJoinApplicable(left, right, leftJoinConditions, leftJoinKeys, rightJoinKeys, joinAlgo)){
+            auto cost = ComputeJoinStats(*left->Stats, *right->Stats, leftJoinKeys, rightJoinKeys, joinAlgo, ctx).Cost;
+            if (bestJoinValid){
+                if (cost < bestCost) {
+                    bestCost = cost;
+                    bestAlgo = joinAlgo;
+                    bestJoinLeftRight = true;
                 }
             } else {
-                res = p1;
+                bestCost = cost;
+                bestAlgo = joinAlgo;
+                bestJoinLeftRight = true;
+                bestJoinValid = true;
             }
         }
-        if (p2) {
-            if (res) {
-                if (p2->Stats->Cost < res->Stats->Cost) {
-                    res = p2;
+
+        if (ctx.IsJoinApplicable(right, left, rightJoinConditions, rightJoinKeys, leftJoinKeys, joinAlgo)){
+            auto cost = ComputeJoinStats(*right->Stats, *left->Stats,  rightJoinKeys, leftJoinKeys, joinAlgo, ctx).Cost;
+            if (bestJoinValid){
+                if (cost < bestCost) {
+                    bestCost = cost;
+                    bestAlgo = joinAlgo;
+                    bestJoinLeftRight = false;
                 }
             } else {
-                res = p2;
+                bestCost = cost;
+                bestAlgo = joinAlgo;
+                bestJoinLeftRight = false;
+                bestJoinValid = true;
             }
         }
     }
 
-    Y_ENSURE(res,"No join was chosen!");
-    return res;
+    Y_ENSURE(bestJoinValid,"No join was chosen!");
+
+    if (bestJoinLeftRight) {
+        return MakeJoinInternal(left, right, leftJoinConditions, leftJoinKeys, rightJoinKeys, EJoinKind::InnerJoin, bestAlgo, ctx); 
+    } else {
+        return MakeJoinInternal(right, left, rightJoinConditions, rightJoinKeys, leftJoinKeys, EJoinKind::InnerJoin, bestAlgo, ctx); 
+    }
 }
 
 /**
@@ -156,30 +319,33 @@ std::shared_ptr<TJoinOptimizerNode> PickBestJoin(std::shared_ptr<IBaseOptimizerN
 std::shared_ptr<TJoinOptimizerNode> PickBestNonReorderabeJoin(std::shared_ptr<IBaseOptimizerNode> left, 
     std::shared_ptr<IBaseOptimizerNode> right, 
     const std::set<std::pair<TJoinColumn, TJoinColumn>>& leftJoinConditions,
+    const TVector<TString>& leftJoinKeys,
+    const TVector<TString>& rightJoinKeys,
     EJoinKind joinKind,
     IProviderContext& ctx) {
 
-    auto res = std::shared_ptr<TJoinOptimizerNode>();
+    EJoinAlgoType bestJoinAlgo;
+    bool bestJoinValid = false;
+    double bestJoinCost;
 
     for ( auto joinAlgo : AllJoinAlgos ) {
-        auto p = ctx.IsJoinApplicable(left, right, leftJoinConditions, joinAlgo) ? 
-            MakeJoin(left, right, leftJoinConditions, joinKind, joinAlgo, ctx) :
-            std::shared_ptr<TJoinOptimizerNode>();
-            
-        if (p) {
-            if (res) {
-                if (p->Stats->Cost < res->Stats->Cost) {
-                    res = p;
+        if (ctx.IsJoinApplicable(left, right, leftJoinConditions, leftJoinKeys, rightJoinKeys, joinAlgo)){
+            auto cost = ComputeJoinStats(*right->Stats, *left->Stats,  rightJoinKeys, leftJoinKeys, joinAlgo, ctx).Cost;
+            if (bestJoinValid) {
+                if (cost < bestJoinCost) {
+                    bestJoinAlgo = joinAlgo;
+                    bestJoinCost = cost;
                 }
             } else {
-                res = p;
+                bestJoinAlgo = joinAlgo;
+                bestJoinCost = cost;
+                bestJoinValid = true;
             }
         }
-
     }
 
-    Y_ENSURE(res,"No join was chosen!");
-    return res;
+    Y_ENSURE(bestJoinValid,"No join was chosen!");
+    return MakeJoin(left, right, leftJoinConditions, leftJoinKeys, rightJoinKeys, joinKind, bestJoinAlgo, true, ctx); 
 }
 
 struct pair_hash {
@@ -273,7 +439,7 @@ struct TGraph {
 
     // Find an edge that connects two subsets of graph's nodes
     // We are guaranteed to find a match
-    TEdge FindCrossingEdge(const std::bitset<N>& S1, const std::bitset<N>& S2) {
+    const TEdge& FindCrossingEdge(const std::bitset<N>& S1, const std::bitset<N>& S2) {
         for(int i = 0; i < NNodes; i++){
             if (!S1[i]) {
                 continue;
@@ -290,7 +456,7 @@ struct TGraph {
             }
         }
         Y_ENSURE(false,"Connecting edge not found!");
-        return TEdge(-1,-1);
+        return TEdge::ErrorEdge;
     }
 
     /**
@@ -341,6 +507,10 @@ struct TGraph {
                         Y_ABORT_UNLESS(maybeEdge1 != Edges.end() && maybeEdge2 != Edges.end());
                         maybeEdge1->JoinConditions.emplace(left, right);
                         maybeEdge2->JoinConditions.emplace(right, left);
+                        maybeEdge1->LeftJoinKeys.emplace_back(left.AttributeName);
+                        maybeEdge1->RightJoinKeys.emplace_back(right.AttributeName);
+                        maybeEdge2->LeftJoinKeys.emplace_back(right.AttributeName);
+                        maybeEdge2->RightJoinKeys.emplace_back(left.AttributeName);
                     }
                 }
             }
@@ -392,7 +562,7 @@ public:
     }
 
     // Run DPccp algorithm and produce the join tree in CBO's internal representation
-    std::shared_ptr<TJoinOptimizerNode> Solve();
+    std::shared_ptr<TJoinOptimizerNodeInternal> Solve();
 
     // Calculate the size of a dynamic programming table with a budget
     ui32 CountCC(ui32 budget);
@@ -490,7 +660,7 @@ template<int N> std::bitset<N> TDPccpSolver<N>::Neighbors(const std::bitset<N>&
 }
 
 // Run the entire DPccp algorithm and compute the optimal join tree
-template<int N> std::shared_ptr<TJoinOptimizerNode> TDPccpSolver<N>::Solve() 
+template<int N> std::shared_ptr<TJoinOptimizerNodeInternal> TDPccpSolver<N>::Solve() 
 {
     // Process singleton sets
     for (int i = NNodes-1; i >= 0; i--) {
@@ -515,7 +685,7 @@ template<int N> std::shared_ptr<TJoinOptimizerNode> TDPccpSolver<N>::Solve()
     }
 
     Y_ENSURE(DpTable.contains(V), "Final relset not in dptable");
-    return std::static_pointer_cast<TJoinOptimizerNode>(DpTable[V]); 
+    return std::static_pointer_cast<TJoinOptimizerNodeInternal>(DpTable[V]); 
 }
 
 /**
@@ -629,9 +799,9 @@ template <int N> void TDPccpSolver<N>::EmitCsgCmp(const std::bitset<N>& S1, cons
 
     std::bitset<N> joined = S1 | S2;
 
-    TEdge e1 = Graph.FindCrossingEdge(S1, S2);
-    TEdge e2 = Graph.FindCrossingEdge(S2, S1);
-    auto bestJoin = PickBestJoin(DpTable[S1], DpTable[S2], e1.JoinConditions, e2.JoinConditions, Pctx);
+    const TEdge& e1 = Graph.FindCrossingEdge(S1, S2);
+    const TEdge& e2 = Graph.FindCrossingEdge(S2, S1);
+    auto bestJoin = PickBestJoin(DpTable[S1], DpTable[S2], e1.JoinConditions, e2.JoinConditions, e1.LeftJoinKeys, e1.RightJoinKeys, Pctx);
 
     if (! DpTable.contains(joined)) {
         DpTable[joined] = bestJoin;
@@ -855,6 +1025,9 @@ TExprBase RearrangeEquiJoinTree(TExprContext& ctx, const TCoEquiJoin& equiJoin,
         .Done();
 }
 
+/**
+ * Collects EquiJoin inputs with statistics for cost based optimization
+*/
 bool DqCollectJoinRelationsWithStats(
     TVector<std::shared_ptr<TRelOptimizerNode>>& rels,
     TTypeAnnotationContext& typesCtx, 
@@ -923,6 +1096,7 @@ std::shared_ptr<TJoinOptimizerNode> ConvertToJoinTree(const TCoEquiJoinTuple& jo
     }
 
     std::set<std::pair<TJoinColumn, TJoinColumn>> joinConds;
+
     size_t joinKeysCount = joinTuple.LeftKeys().Size() / 2;
     for (size_t i = 0; i < joinKeysCount; ++i) {
         size_t keyIndex = i * 2;
@@ -1001,7 +1175,8 @@ void ComputeStatistics(const std::shared_ptr<TJoinOptimizerNode>& join, IProvide
     if (join->RightArg->Kind == EOptimizerNodeKind::JoinNodeType) {
         ComputeStatistics(static_pointer_cast<TJoinOptimizerNode>(join->RightArg), ctx);
     }
-    join->Stats = std::make_shared<TOptimizerStatistics>(ComputeJoinStats(*join->LeftArg->Stats, *join->RightArg->Stats, join->JoinConditions, EJoinAlgoType::DictJoin, ctx));
+    join->Stats = std::make_shared<TOptimizerStatistics>(ComputeJoinStats(*join->LeftArg->Stats, *join->RightArg->Stats, 
+        join->LeftJoinKeys, join->RightJoinKeys, EJoinAlgoType::DictJoin, ctx));
 }
 
 /**
@@ -1011,15 +1186,16 @@ void ComputeStatistics(const std::shared_ptr<TJoinOptimizerNode>& join, IProvide
 */
 std::shared_ptr<TJoinOptimizerNode> OptimizeSubtree(const std::shared_ptr<TJoinOptimizerNode>& joinTree, ui32 maxDPccpDPTableSize, IProviderContext& ctx) {
     if (!joinTree->IsReorderable) {
-        return PickBestNonReorderabeJoin(joinTree->LeftArg, joinTree->RightArg, joinTree->JoinConditions, joinTree->JoinType, ctx);
+        return PickBestNonReorderabeJoin(joinTree->LeftArg, joinTree->RightArg, joinTree->JoinConditions, 
+            joinTree->LeftJoinKeys, joinTree->RightJoinKeys, joinTree->JoinType, ctx);
     }
 
-    TGraph<128> joinGraph;
     TVector<std::shared_ptr<IBaseOptimizerNode>> rels;
     std::set<std::pair<TJoinColumn, TJoinColumn>> joinConditions;
-
     ExtractRelsAndJoinConditions(joinTree, rels, joinConditions);
 
+    TGraph<128> joinGraph;
+
     for (size_t i = 0; i < rels.size(); i++) {
         joinGraph.AddNode(i, rels[i]->Labels());
     }
@@ -1029,6 +1205,7 @@ std::shared_ptr<TJoinOptimizerNode> OptimizeSubtree(const std::shared_ptr<TJoinO
     // computed statistics
     if (rels.size() >= 128) {
         ComputeStatistics(joinTree, ctx);
+        YQL_CLOG(TRACE, CoreDq) << "Too many rels";
         return joinTree;
     }
 
@@ -1065,7 +1242,7 @@ std::shared_ptr<TJoinOptimizerNode> OptimizeSubtree(const std::shared_ptr<TJoinO
     }
 
     // feed the graph to DPccp algorithm
-    std::shared_ptr<TJoinOptimizerNode> result = solver.Solve();
+    auto result = solver.Solve();
 
     if (NYql::NLog::YqlLogger().NeedToLog(NYql::NLog::EComponent::ProviderKqp, NYql::NLog::ELevel::TRACE)) {
         std::stringstream str;
@@ -1074,7 +1251,7 @@ std::shared_ptr<TJoinOptimizerNode> OptimizeSubtree(const std::shared_ptr<TJoinO
         YQL_CLOG(TRACE, CoreDq) << str.str();
     }
 
-    return result;
+    return ConvertFromInternal(result);
 }
 
 class TOptimizerNativeNew: public IOptimizerNew {
@@ -1083,6 +1260,7 @@ public:
         : IOptimizerNew(ctx), MaxDPccpDPTableSize(maxDPccpDPTableSize) { }
 
     std::shared_ptr<TJoinOptimizerNode>  JoinSearch(const std::shared_ptr<TJoinOptimizerNode>& joinTree) override {
+
         // Traverse the join tree and generate a list of non-orderable joins in a post-order
         TVector<std::shared_ptr<TJoinOptimizerNode>> nonOrderables;
         ExtractNonOrderables(joinTree, nonOrderables);
@@ -1095,7 +1273,8 @@ public:
             if (join->RightArg->Kind == EOptimizerNodeKind::JoinNodeType) {
                 join->RightArg = OptimizeSubtree(static_pointer_cast<TJoinOptimizerNode>(join->RightArg), MaxDPccpDPTableSize, Pctx);
             }  
-            join->Stats = std::make_shared<TOptimizerStatistics>(ComputeJoinStats(*join->LeftArg->Stats, *join->RightArg->Stats, join->JoinConditions, EJoinAlgoType::DictJoin, Pctx));
+            join->Stats = std::make_shared<TOptimizerStatistics>(ComputeJoinStats(*join->LeftArg->Stats, *join->RightArg->Stats, 
+                join->LeftJoinKeys, join->RightJoinKeys, EJoinAlgoType::DictJoin, Pctx));
         }
 
         // Optimize the root
@@ -1152,6 +1331,13 @@ TExprBase DqOptimizeEquiJoinWithCosts(const TExprBase& node, TExprContext& ctx,
     // Generate an initial tree
     auto joinTree = ConvertToJoinTree(joinTuple,rels);
 
+    if (NYql::NLog::YqlLogger().NeedToLog(NYql::NLog::EComponent::ProviderKqp, NYql::NLog::ELevel::TRACE)) {
+        std::stringstream str;
+        str << "Converted join tree:\n";
+        joinTree->Print(str);
+        YQL_CLOG(TRACE, CoreDq) << str.str();
+    }
+
     auto opt = TOptimizerNativeNew(providerCtx, maxDPccpDPTableSize);
     joinTree = opt.JoinSearch(joinTree);
 
@@ -1173,7 +1359,7 @@ public:
     TOutput JoinSearch() override {
         auto dummyProviderCtx = TDummyProviderContext();
         TDPccpSolver<128> solver(JoinGraph, Rels, dummyProviderCtx);
-        std::shared_ptr<TJoinOptimizerNode> result = solver.Solve();
+        std::shared_ptr<TJoinOptimizerNode> result = ConvertFromInternal(solver.Solve());
         if (Log) {
             std::stringstream str;
             str << "Join tree after cost based optimization:\n";