Add support for KqpOlapApply in explain (#18431)

Author: denis0x0D

ydb/core/kqp/opt/kqp_query_plan.cpp

@@ -1179,6 +1179,8 @@ class TxPlanSerializer {
                     return Sprintf(strRegexp[compSign].c_str(), attr.c_str(), value.c_str());
                 }
             }
+        } else if (auto olapApply = TMaybeNode<TKqpOlapApply>(node)) {
+            return NPlanUtils::ExtractPredicate(olapApply.Cast().Lambda()).Body;
         }
 
         for (const auto& child: node->Children()) {

ydb/core/kqp/opt/kqp_statistics_transformer.cpp

@@ -478,6 +478,8 @@ class TKqpOlapPredicateSelectivityComputer: public TPredicateSelectivityComputer
                 tmpSelectivity *= Compute(andNode.Cast().Arg(i));
             }
             resSelectivity = tmpSelectivity;
+        } else if (auto olapApply = input.Maybe<TKqpOlapApply>()) {
+            resSelectivity = TPredicateSelectivityComputer::Compute(olapApply.Cast().Lambda().Body());
         } else if (auto orNode = input.Maybe<TKqpOlapOr>()) {
             double tmpSelectivity = 0.0;
             for (size_t i = 0; i < orNode.Cast().ArgCount(); i++) {

ydb/core/kqp/ut/join/kqp_join_order_ut.cpp

@@ -476,10 +476,9 @@ Y_UNIT_TEST_SUITE(OlapEstimationRowsCorrectness) {
         TestOlapEstimationRowsCorrectness("queries/tpch2.sql", "stats/tpch1000s.json");
     }
 
-    // FIXME: Cardinality estimation is broken because of new type of OLAP pushdown
-    // Y_UNIT_TEST(TPCH3) {
-    //    TestOlapEstimationRowsCorrectness("queries/tpch3.sql", "stats/tpch1000s.json");
-    // }
+    Y_UNIT_TEST(TPCH3) {
+       TestOlapEstimationRowsCorrectness("queries/tpch3.sql", "stats/tpch1000s.json");
+    }
 
     Y_UNIT_TEST(TPCH5) {
         TestOlapEstimationRowsCorrectness("queries/tpch5.sql", "stats/tpch1000s.json");

ydb/library/yql/dq/opt/dq_opt_predicate_selectivity.cpp

@@ -24,9 +24,9 @@ namespace {
      * Check if a callable is an attribute of some table
      * Currently just return a boolean and cover only basic cases
      */
-    TMaybe<TCoMember> IsAttribute(const TExprBase& input) {
+    std::optional<TString> IsAttribute(const TExprBase& input) {
         if (auto member = input.Maybe<TCoMember>()) {
-            return member.Cast();
+            return TString(member.Cast().Name());
         } else if (auto cast = input.Maybe<TCoSafeCast>()) {
             return IsAttribute(cast.Cast().Value());
         } else if (auto ifPresent = input.Maybe<TCoIfPresent>()) {
@@ -42,8 +42,37 @@ namespace {
         } else if (auto exists = input.Maybe<TCoExists>()) {
             auto child = TExprBase(input.Ptr()->ChildRef(0));
             return IsAttribute(child);
+        } else if (auto argument = input.Maybe<TCoArgument>()) {
+            TString argumentName = TString(argument.Cast().Name());
+            TStringBuf olapApplyMemberPrefix = "members_";
+            if (argumentName.StartsWith(olapApplyMemberPrefix)) {
+                return argumentName.substr(olapApplyMemberPrefix.length(), argumentName.size() - olapApplyMemberPrefix.length());
+            } else {
+                return argumentName;
+            }
         }
+        return std::nullopt;
+    }
 
+    TMaybe<TCoMember> IsMember(const TExprBase& input) {
+        if (auto member = input.Maybe<TCoMember>()) {
+            return member.Cast();
+        } else if (auto cast = input.Maybe<TCoSafeCast>()) {
+            return IsMember(cast.Cast().Value());
+        } else if (auto ifPresent = input.Maybe<TCoIfPresent>()) {
+            return IsMember(ifPresent.Cast().Optional());
+        } else if (auto just = input.Maybe<TCoJust>()) {
+            return IsMember(just.Cast().Input());
+        } else if (input.Ptr()->IsCallable("PgCast")) {
+            auto child = TExprBase(input.Ptr()->ChildRef(0));
+            return IsMember(child);
+        } else if (input.Ptr()->IsCallable("FromPg")) {
+            auto child = TExprBase(input.Ptr()->ChildRef(0));
+            return IsMember(child);
+        } else if (auto exists = input.Maybe<TCoExists>()) {
+            auto child = TExprBase(input.Ptr()->ChildRef(0));
+            return IsMember(child);
+        }
         return Nothing();
     }
 
@@ -223,12 +252,12 @@ double NYql::NDq::TPredicateSelectivityComputer::ComputeInequalitySelectivity(co
         return ComputeInequalitySelectivity(right, left, GetOppositePredicateType(predicate));
     }
 
-    if (auto maybeMember = IsAttribute(left)) {
+    if (auto attribute = IsAttribute(left)) {
         // It seems like this is not possible in current version.
         if (IsAttribute(right)) {
             return 0.3;
         } else if (IsConstantExprWithParams(right.Ptr())) {
-            const TString attributeName = maybeMember.Get()->Name().StringValue();
+            const TString attributeName = attribute.value();
             if (!IsConstantExpr(right.Ptr())) {
                 return DefaultSelectivity(Stats, attributeName);
             }
@@ -259,26 +288,32 @@ double NYql::NDq::TPredicateSelectivityComputer::ComputeEqualitySelectivity(cons
         return ComputeEqualitySelectivity(right, left);
     }
 
-    if (auto maybeMember = IsAttribute(left)) {
+    if (auto attribute = IsAttribute(left)) {
         // In case both arguments refer to an attribute, return 0.2
-        if (auto maybeAnotherMember = IsAttribute(right)) {
+        if (IsAttribute(right)) {
             if (CollectMemberEqualities) {
-                MemberEqualities.Add(*maybeMember.Get(), *maybeAnotherMember.Get());
+                auto maybeMember = IsMember(left);
+                auto maybeAnotherMember = IsMember(right);
+                if (maybeMember && maybeAnotherMember) {
+                    MemberEqualities.Add(*maybeMember.Get(), *maybeAnotherMember.Get());
+                }
             }
             return 0.3;
         }
         // In case the right side is a constant that can be extracted, compute the selectivity using statistics
         // Currently, with the basic statistics we just return 1/nRows
 
         else if (IsConstantExprWithParams(right.Ptr())) {
-            TString attributeName = maybeMember.Get()->Name().StringValue();
+            TString attributeName = attribute.value();
             if (!IsConstantExpr(right.Ptr())) {
                 return DefaultSelectivity(Stats, attributeName);
             }
 
             if (Stats == nullptr || Stats->ColumnStatistics == nullptr) {
                 if (CollectColumnsStatUsedMembers) {
-                    ColumnStatsUsedMembers.AddEquality(*maybeMember.Get());
+                    if (auto maybeMember = IsMember(left)) {
+                        ColumnStatsUsedMembers.AddEquality(*maybeMember.Get());
+                    }
                 }
                 return DefaultSelectivity(Stats, attributeName);
             }

ydb/library/yql/utils/plan/plan_utils.cpp

@@ -285,6 +285,13 @@ TString PrettyExprStr(const TExprBase& expr) {
 
         return TStringBuilder() << "[" << JoinStrings(std::move(items), ",") << "]";
     } else {
+        if (auto arg = expr.Maybe<TCoArgument>()) {
+            TStringBuf yqlOlapApplyMember = "members_";
+            TString argumentName = TString(arg.Cast().Name());
+            if (argumentName.StartsWith(yqlOlapApplyMember)) {
+                return argumentName.substr(yqlOlapApplyMember.length(), argumentName.length() - yqlOlapApplyMember.length());
+            }
+        }
         auto raw = TString(expr.Ref().Content());
         // return raw.StartsWith("_yql_agg_") ? "" : raw;
         return raw;