Added statistics inference for DqJoins (#11994)

pavelvelikhov · web-flow · commit 806d77103e5c · 2024-11-26T16:59:51.000+03:00
diff --git a/ydb/core/kqp/ut/join/data/join_order/tpcds64_1000s.json b/ydb/core/kqp/ut/join/data/join_order/tpcds64_1000s.json
diff --git a/ydb/core/kqp/ut/join/data/join_order/tpcds64_1000s_column_store.json b/ydb/core/kqp/ut/join/data/join_order/tpcds64_1000s_column_store.json
diff --git a/ydb/core/kqp/ut/join/data/join_order/tpch2_1000s.json b/ydb/core/kqp/ut/join/data/join_order/tpch2_1000s.json
@@ -7,39 +7,39 @@
         "args":
           [
             {
-              "op_name":"InnerJoin (Grace)",
+              "op_name":"TableFullScan",
+              "table":"partsupp"
+            },
+            {
+              "op_name":"InnerJoin (MapJoin)",
               "args":
                 [
                   {
                     "op_name":"TableFullScan",
-                    "table":"partsupp"
+                    "table":"supplier"
                   },
                   {
                     "op_name":"InnerJoin (MapJoin)",
                     "args":
                       [
                         {
-                          "op_name":"TableFullScan",
-                          "table":"supplier"
+                          "op_name":"TableLookup",
+                          "table":"region"
                         },
                         {
-                          "op_name":"InnerJoin (MapJoin)",
-                          "args":
-                            [
-                              {
-                                "op_name":"TableLookup",
-                                "table":"region"
-                              },
-                              {
-                                "op_name":"TableFullScan",
-                                "table":"nation"
-                              }
-                            ]
+                          "op_name":"TableFullScan",
+                          "table":"nation"
                         }
                       ]
                   }
                 ]
-            },
+            }
+          ]
+      },
+      {
+        "op_name":"InnerJoin (Grace)",
+        "args":
+          [
             {
               "op_name":"InnerJoin (Grace)",
               "args":
@@ -73,12 +73,12 @@
                       ]
                   }
                 ]
+            },
+            {
+              "op_name":"TableFullScan",
+              "table":"part"
             }
           ]
-      },
-      {
-        "op_name":"TableFullScan",
-        "table":"part"
       }
     ]
 }
diff --git a/ydb/core/kqp/ut/join/data/join_order/tpch2_1000s_column_store.json b/ydb/core/kqp/ut/join/data/join_order/tpch2_1000s_column_store.json
@@ -7,39 +7,39 @@
         "args":
           [
             {
-              "op_name":"InnerJoin (Grace)",
+              "op_name":"TableFullScan",
+              "table":"partsupp"
+            },
+            {
+              "op_name":"InnerJoin (MapJoin)",
               "args":
                 [
                   {
                     "op_name":"TableFullScan",
-                    "table":"partsupp"
+                    "table":"supplier"
                   },
                   {
                     "op_name":"InnerJoin (MapJoin)",
                     "args":
                       [
                         {
                           "op_name":"TableFullScan",
-                          "table":"supplier"
+                          "table":"nation"
                         },
                         {
-                          "op_name":"InnerJoin (MapJoin)",
-                          "args":
-                            [
-                              {
-                                "op_name":"TableFullScan",
-                                "table":"nation"
-                              },
-                              {
-                                "op_name":"TableFullScan",
-                                "table":"region"
-                              }
-                            ]
+                          "op_name":"TableFullScan",
+                          "table":"region"
                         }
                       ]
                   }
                 ]
-            },
+            }
+          ]
+      },
+      {
+        "op_name":"InnerJoin (Grace)",
+        "args":
+          [
             {
               "op_name":"InnerJoin (Grace)",
               "args":
@@ -73,12 +73,12 @@
                       ]
                   }
                 ]
+            },
+            {
+              "op_name":"TableFullScan",
+              "table":"part"
             }
           ]
-      },
-      {
-        "op_name":"TableFullScan",
-        "table":"part"
       }
     ]
 }
diff --git a/ydb/library/yql/dq/opt/dq_opt_stat.cpp b/ydb/library/yql/dq/opt/dq_opt_stat.cpp
@@ -369,6 +369,69 @@ void InferStatisticsForGraceJoin(const TExprNode::TPtr& input, TTypeAnnotationCo
     YQL_CLOG(TRACE, CoreDq) << "Infer statistics for GraceJoin: " << resStats->ToString();
 }
 
+/**
+ * Infer statistics for DqJoin
+ * DqJoin is an intermediary join representantation in Dq
+ */
+void InferStatisticsForDqJoin(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx, const IProviderContext& ctx, TCardinalityHints hints) {
+    auto inputNode = TExprBase(input);
+    auto join = inputNode.Cast<TDqJoin>();
+
+    auto leftArg = join.LeftInput();
+    auto rightArg = join.RightInput();
+
+    auto leftStats = typeCtx->GetStats(leftArg.Raw());
+    auto rightStats = typeCtx->GetStats(rightArg.Raw());
+
+    if (!leftStats || !rightStats) {
+        return;
+    }
+
+    auto joinAlgo = FromString<EJoinAlgoType>(join.JoinAlgo().StringValue());
+    if (joinAlgo == EJoinAlgoType::Undefined) {
+        return;
+    }
+
+    auto leftLabels = InferLabels(leftStats, join.LeftJoinKeyNames());
+    auto rightLabels = InferLabels(rightStats, join.RightJoinKeyNames());
+
+    leftStats = ApplyCardinalityHints(leftStats, leftLabels, hints);
+    rightStats = ApplyCardinalityHints(rightStats, rightLabels, hints);
+
+    TVector<TJoinColumn> leftJoinKeys;
+    TVector<TJoinColumn> rightJoinKeys;
+
+    for (size_t i=0; i<join.LeftJoinKeyNames().Size(); i++) {
+        auto alias = ExtractAlias(join.LeftJoinKeyNames().Item(i).StringValue());
+        auto attrName = RemoveAliases(join.LeftJoinKeyNames().Item(i).StringValue());
+        leftJoinKeys.push_back(TJoinColumn(alias, attrName));
+    }
+    for (size_t i=0; i<join.RightJoinKeyNames().Size(); i++) {
+        auto alias = ExtractAlias(join.RightJoinKeyNames().Item(i).StringValue());
+        auto attrName = RemoveAliases(join.RightJoinKeyNames().Item(i).StringValue());
+        rightJoinKeys.push_back(TJoinColumn(alias, attrName));
+    }
+
+    auto unionOfLabels = UnionLabels(leftLabels, rightLabels);
+
+    auto resStats = std::make_shared<TOptimizerStatistics>(
+            ctx.ComputeJoinStats(
+                *leftStats,
+                *rightStats,
+                leftJoinKeys,
+                rightJoinKeys, 
+                joinAlgo,
+                ConvertToJoinKind(join.JoinType().StringValue()),
+                FindCardHint(unionOfLabels, hints)
+            )
+        );
+
+    resStats->Labels = std::make_shared<TVector<TString>>();
+    resStats->Labels->insert(resStats->Labels->begin(), unionOfLabels.begin(), unionOfLabels.end());
+    typeCtx->SetStats(join.Raw(), resStats);
+    YQL_CLOG(TRACE, CoreDq) << "Infer statistics for DqJoin: " << resStats->ToString();
+}
+
 /**
  * Infer statistics for DqSource
  *
diff --git a/ydb/library/yql/dq/opt/dq_opt_stat.h b/ydb/library/yql/dq/opt/dq_opt_stat.h
@@ -19,6 +19,7 @@ void InferStatisticsForStage(const TExprNode::TPtr& input, TTypeAnnotationContex
 void InferStatisticsForDqSource(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 void InferStatisticsForGraceJoin(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx, const IProviderContext& ctx, TCardinalityHints hints = {});
 void InferStatisticsForMapJoin(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx, const IProviderContext& ctx, TCardinalityHints hints = {});
+void InferStatisticsForDqJoin(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx, const IProviderContext& ctx, TCardinalityHints hints = {});
 void InferStatisticsForAsList(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 bool InferStatisticsForListParam(const TExprNode::TPtr& input, TTypeAnnotationContext* typeCtx);
 
diff --git a/ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp b/ydb/library/yql/dq/opt/dq_opt_stat_transformer_base.cpp
@@ -65,6 +65,9 @@ bool TDqStatisticsTransformerBase::BeforeLambdas(const TExprNode::TPtr& input, T
     else if(TCoGraceJoinCore::Match(input.Get())) {
         InferStatisticsForGraceJoin(input, TypeCtx, Pctx, CardinalityHints);
     }
+    else if (TDqJoin::Match(input.Get())) {
+        InferStatisticsForDqJoin(input, TypeCtx, Pctx, CardinalityHints);
+    }
 
     // Do nothing in case of EquiJoin, otherwise the EquiJoin rule won't fire
     else if(TCoEquiJoin::Match(input.Get())){

Original file line number	Diff line number	Diff line change
`@@ -7,39 +7,39 @@`
`7`	`7`	`"args":`
`8`	`8`	`[`
`9`	`9`	`{`
`10`		`- "op_name":"InnerJoin (Grace)",`
	`10`	`+ "op_name":"TableFullScan",`
	`11`	`+ "table":"partsupp"`
	`12`	`+ },`
	`13`	`+ {`
	`14`	`+ "op_name":"InnerJoin (MapJoin)",`
`11`	`15`	`"args":`
`12`	`16`	`[`
`13`	`17`	`{`
`14`	`18`	`"op_name":"TableFullScan",`
`15`		`- "table":"partsupp"`
	`19`	`+ "table":"supplier"`
`16`	`20`	`},`
`17`	`21`	`{`
`18`	`22`	`"op_name":"InnerJoin (MapJoin)",`
`19`	`23`	`"args":`
`20`	`24`	`[`
`21`	`25`	`{`
`22`		`- "op_name":"TableFullScan",`
`23`		`- "table":"supplier"`
	`26`	`+ "op_name":"TableLookup",`
	`27`	`+ "table":"region"`
`24`	`28`	`},`
`25`	`29`	`{`
`26`		`- "op_name":"InnerJoin (MapJoin)",`
`27`		`- "args":`
`28`		`- [`
`29`		`- {`
`30`		`- "op_name":"TableLookup",`
`31`		`- "table":"region"`
`32`		`- },`
`33`		`- {`
`34`		`- "op_name":"TableFullScan",`
`35`		`- "table":"nation"`
`36`		`- }`
`37`		`- ]`
	`30`	`+ "op_name":"TableFullScan",`
	`31`	`+ "table":"nation"`
`38`	`32`	`}`
`39`	`33`	`]`
`40`	`34`	`}`
`41`	`35`	`]`
`42`		`- },`
	`36`	`+ }`
	`37`	`+ ]`
	`38`	`+ },`
	`39`	`+ {`
	`40`	`+ "op_name":"InnerJoin (Grace)",`
	`41`	`+ "args":`
	`42`	`+ [`
`43`	`43`	`{`
`44`	`44`	`"op_name":"InnerJoin (Grace)",`
`45`	`45`	`"args":`
`@@ -73,12 +73,12 @@`
`73`	`73`	`]`
`74`	`74`	`}`
`75`	`75`	`]`
	`76`	`+ },`
	`77`	`+ {`
	`78`	`+ "op_name":"TableFullScan",`
	`79`	`+ "table":"part"`
`76`	`80`	`}`
`77`	`81`	`]`
`78`		`- },`
`79`		`- {`
`80`		`- "op_name":"TableFullScan",`
`81`		`- "table":"part"`
`82`	`82`	`}`
`83`	`83`	`]`
`84`	`84`	`}`
Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,9 @@ bool TDqStatisticsTransformerBase::BeforeLambdas(const TExprNode::TPtr& input, T`
`65`	`65`	`else if(TCoGraceJoinCore::Match(input.Get())) {`
`66`	`66`	`InferStatisticsForGraceJoin(input, TypeCtx, Pctx, CardinalityHints);`
`67`	`67`	`}`
	`68`	`+ else if (TDqJoin::Match(input.Get())) {`
	`69`	`+ InferStatisticsForDqJoin(input, TypeCtx, Pctx, CardinalityHints);`
	`70`	`+ }`
`68`	`71`
`69`	`72`	`// Do nothing in case of EquiJoin, otherwise the EquiJoin rule won't fire`
`70`	`73`	`else if(TCoEquiJoin::Match(input.Get())){`