Merge #14222 (#14274)

Author: vitalyisaev2

ydb/library/yql/providers/common/pushdown/collection.cpp

@@ -9,7 +9,7 @@ namespace NYql::NPushdown {
 
 // Collects subpredicate that we can then push down
 void CollectPredicates(const NNodes::TExprBase& predicate, TPredicateNode& predicateTree,
-                       const TExprNode* lambdaArg, const NNodes::TExprBase& lambdaBody,
+                       const NNodes::TExprBase& lambdaArg, const NNodes::TExprBase& lambdaBody,
                        const TSettings& settings);
 
 } // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/collection.h

@@ -0,0 +1,80 @@
+#include "predicate_node.h"
+
+#include <ydb/library/yql/utils/log/log.h>
+#include <ydb/library/yql/core/expr_nodes/yql_expr_nodes.h>
+#include <ydb/library/yql/providers/common/provider/yql_provider.h>
+#include <ydb/library/yql/providers/common/pushdown/collection.h>
+
+namespace NYql::NPushdown {
+
+using namespace NNodes;
+
+namespace {
+
+TPredicateNode SplitForPartialPushdown(const NPushdown::TPredicateNode& predicateTree, TExprContext& ctx, TPositionHandle pos, const TSettings& settings) {
+    if (predicateTree.CanBePushed) {
+        return predicateTree;
+    }
+
+    if (predicateTree.Op != NPushdown::EBoolOp::And && (!settings.IsEnabled(TSettings::EFeatureFlag::SplitOrOperator) || predicateTree.Op != NPushdown::EBoolOp::Or)) {
+        // Predicate can't be split, so return invalid value and skip this branch
+        return NPushdown::TPredicateNode();
+    }
+
+    std::vector<NPushdown::TPredicateNode> pushable;
+    for (auto& predicate : predicateTree.Children) {
+        NPushdown::TPredicateNode pushablePredicate = SplitForPartialPushdown(predicate, ctx, pos, settings);
+        if (pushablePredicate.IsValid()) {
+            pushable.emplace_back(pushablePredicate);
+        } else if (predicateTree.Op == NPushdown::EBoolOp::Or) {
+            // One of the OR branch was invalid, so the whole predicate is invalid
+            return NPushdown::TPredicateNode();
+        }
+    }
+
+    NPushdown::TPredicateNode predicateToPush;
+    predicateToPush.SetPredicates(pushable, ctx, pos, predicateTree.Op);
+    return predicateToPush;
+}
+
+}
+
+NPushdown::TPredicateNode MakePushdownNode(const NNodes::TCoLambda& lambda, TExprContext& ctx, const TPositionHandle& pos, const TSettings& settings) {
+    auto lambdaArg = lambda.Args().Arg(0).Ptr();
+
+    YQL_LOG(TRACE) << "Push filter. Initial filter lambda: " << NCommon::ExprToPrettyString(ctx, lambda.Ref());
+
+    auto maybeOptionalIf = lambda.Body().Maybe<TCoOptionalIf>();
+    if (!maybeOptionalIf.IsValid()) { // Nothing to push
+        return {};
+    }
+
+    TCoOptionalIf optionalIf = maybeOptionalIf.Cast();
+    NPushdown::TPredicateNode predicateTree(optionalIf.Predicate());
+    NPushdown::CollectPredicates(optionalIf.Predicate(), predicateTree, TExprBase(lambdaArg), TExprBase(lambdaArg), settings);
+    YQL_ENSURE(predicateTree.IsValid(), "Collected filter predicates are invalid");
+
+    return SplitForPartialPushdown(predicateTree, ctx, pos, settings);
+}
+
+TMaybeNode<TCoLambda> MakePushdownPredicate(const TCoLambda& lambda, TExprContext& ctx, const TPositionHandle& pos, const TSettings& settings) {
+    NPushdown::TPredicateNode predicateToPush = MakePushdownNode(lambda, ctx, pos, settings);
+    if (!predicateToPush.IsValid()) {
+        return {};
+    }
+
+    // clang-format off
+    auto newFilterLambda = Build<TCoLambda>(ctx, pos)
+        .Args({"filter_row"})
+        .Body<TExprApplier>()
+            .Apply(predicateToPush.ExprNode.Cast())
+            .With(lambda.Args().Arg(0), "filter_row")
+            .Build()
+        .Done();
+    // clang-format on
+
+    YQL_LOG(INFO) << "Push filter lambda: " << NCommon::ExprToPrettyString(ctx, *newFilterLambda.Ptr());
+    return newFilterLambda;
+}
+
+} // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/physical_opt.cpp

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "predicate_node.h"
+
+#include <ydb/library/yql/ast/yql_expr.h>
+#include <ydb/library/yql/ast/yql_pos_handle.h>
+#include <ydb/library/yql/core/expr_nodes/yql_expr_nodes.h>
+#include <ydb/library/yql/core/expr_nodes_gen/yql_expr_nodes_gen.h>
+
+namespace NYql::NPushdown {
+
+NPushdown::TPredicateNode MakePushdownNode(const NNodes::TCoLambda& lambda, TExprContext& ctx, const TPositionHandle& pos, const TSettings& settings);
+NNodes::TMaybeNode<NNodes::TCoLambda> MakePushdownPredicate(const NNodes::TCoLambda& lambda, TExprContext& ctx, const TPositionHandle& pos, const TSettings& settings);
+
+} // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/physical_opt.h

@@ -30,26 +30,48 @@ bool TPredicateNode::IsValid() const {
     return res && ExprNode.IsValid();
 }
 
-void TPredicateNode::SetPredicates(const std::vector<TPredicateNode>& predicates, TExprContext& ctx, TPositionHandle pos) {
+bool TPredicateNode::IsEmpty() const {
+    if (!ExprNode || !IsValid()) {
+        return true;
+    }
+    if (const auto maybeBool = ExprNode.Maybe<NNodes::TCoBool>()) {
+        return TStringBuf(maybeBool.Cast().Literal()) == "true"sv;
+    }
+    return false;
+}
+
+void TPredicateNode::SetPredicates(const std::vector<TPredicateNode>& predicates, TExprContext& ctx, TPositionHandle pos, EBoolOp op) {
     auto predicatesSize = predicates.size();
     if (predicatesSize == 0) {
         return;
-    } else if (predicatesSize == 1) {
+    }
+    if (predicatesSize == 1) {
         *this = predicates[0];
-    } else {
-        Op = EBoolOp::And;
-        Children = predicates;
-        CanBePushed = true;
-
-        TVector<NNodes::TExprBase> exprNodes;
-        exprNodes.reserve(predicatesSize);
-        for (auto& pred : predicates) {
-            exprNodes.emplace_back(pred.ExprNode.Cast());
-            CanBePushed &= pred.CanBePushed;
-        }
-        ExprNode = NNodes::Build<NNodes::TCoAnd>(ctx, pos)
-            .Add(exprNodes)
-            .Done();
+        return;
+    }
+
+    Op = op;
+    Children = predicates;
+    CanBePushed = true;
+
+    TVector<NNodes::TExprBase> exprNodes;
+    exprNodes.reserve(predicatesSize);
+    for (auto& pred : predicates) {
+        exprNodes.emplace_back(pred.ExprNode.Cast());
+        CanBePushed &= pred.CanBePushed;
+    }
+
+    switch (op) {
+        case EBoolOp::And:
+            ExprNode = NNodes::Build<NNodes::TCoAnd>(ctx, pos).Add(exprNodes).Done();
+            break;
+
+        case EBoolOp::Or:
+            ExprNode = NNodes::Build<NNodes::TCoOr>(ctx, pos).Add(exprNodes).Done();
+            break;
+
+        default:
+            throw yexception() << "Unsupported operator for predicate node creation: " << static_cast<int>(op);
     }
 }
 

ydb/library/yql/providers/common/pushdown/predicate_node.cpp

@@ -23,7 +23,8 @@ struct TPredicateNode {
     ~TPredicateNode();
 
     bool IsValid() const;
-    void SetPredicates(const std::vector<TPredicateNode>& predicates, TExprContext& ctx, TPositionHandle pos);
+    bool IsEmpty() const;
+    void SetPredicates(const std::vector<TPredicateNode>& predicates, TExprContext& ctx, TPositionHandle pos, EBoolOp op);
 
     NNodes::TMaybeNode<NNodes::TExprBase> ExprNode;
     std::vector<TPredicateNode> Children;

ydb/library/yql/providers/common/pushdown/predicate_node.h

@@ -10,8 +10,16 @@ void TSettings::Enable(ui64 flagsMask, bool set) {
     }
 }
 
+void TSettings::EnableFunction(const TString& functionName) {
+    EnabledFunctions.insert(functionName);
+}
+
 bool TSettings::IsEnabled(EFeatureFlag flagMask) const {
     return (FeatureFlags & flagMask) != 0;
 }
 
+bool TSettings::IsEnabledFunction(const TString& functionName) const {
+    return EnabledFunctions.contains(functionName);
+}
+
 } // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/settings.cpp

@@ -3,6 +3,8 @@
 
 #include <util/system/types.h>
 
+#include <unordered_set>
+
 namespace NYql::NPushdown {
 
 struct TSettings {
@@ -25,7 +27,20 @@ struct TSettings {
         UnaryOperators = 1 << 15, // -, Abs, Size
         DoNotCheckCompareArgumentsTypes = 1 << 16,
         TimestampCtor = 1 << 17,
-        JustPassthroughOperators = 1 << 18 // if + coalesce + just
+        JustPassthroughOperators = 1 << 18, // if + coalesce + just
+        InOperator = 1 << 19, // IN()
+        IsDistinctOperator = 1 << 20, // IS NOT DISTINCT FROM / IS DISTINCT FROM 
+        DivisionExpressions = 1 << 21, // %, / -- NOTE: division by zero is not handled and also pushdown
+
+        // Option which enables partial pushdown for sequence of OR
+        // For example next predicate:
+        // ($A AND $B) OR ($C AND $D)
+        // May be partially pushdowned as:
+        // $A OR $C
+        // In case of unsupported / complicated expressions $B and $D
+        SplitOrOperator = 1 << 22,
+        ToBytesFromStringExpressions = 1 << 23, // ToBytes(string like)
+        FlatMapOverOptionals = 1 << 24 // FlatMap(Optional<T>, Lmabda (T) -> Optional<U>)
     };
 
     explicit TSettings(NLog::EComponent logComponent)
@@ -37,15 +52,20 @@ struct TSettings {
 
     void Enable(ui64 flagsMask, bool set = true);
 
+    void EnableFunction(const TString& functionName);
+
     bool IsEnabled(EFeatureFlag flagMask) const;
 
+    bool IsEnabledFunction(const TString& functionName) const;
+
     NLog::EComponent GetLogComponent() const {
         return LogComponent;
     }
 
 private:
     const NLog::EComponent LogComponent;
     ui64 FeatureFlags = 0;
+    std::unordered_set<TString> EnabledFunctions;
 };
 
 } // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/settings.h

@@ -0,0 +1,36 @@
+#include "predicate_node.h"
+
+#include <ydb/library/yql/core/expr_nodes/yql_expr_nodes.h>
+#include <ydb/library/yql/core/yql_expr_type_annotation.h>
+
+namespace NYql::NPushdown {
+
+IGraphTransformer::TStatus AnnotateFilterPredicate(const TExprNode::TPtr& input, size_t childIndex, const TStructExprType* itemType, TExprContext& ctx) {
+    if (childIndex >= input->ChildrenSize()) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    auto& filterLambda = input->ChildRef(childIndex);
+    if (!EnsureLambda(*filterLambda, ctx)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (!UpdateLambdaAllArgumentsTypes(filterLambda, {itemType}, ctx)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (const auto* filterLambdaType = filterLambda->GetTypeAnn()) {
+        if (filterLambdaType->GetKind() != ETypeAnnotationKind::Data) {
+            return IGraphTransformer::TStatus::Error;
+        }
+        const TDataExprType* dataExprType = static_cast<const TDataExprType*>(filterLambdaType);
+        if (dataExprType->GetSlot() != EDataSlot::Bool) {
+            return IGraphTransformer::TStatus::Error;
+        }
+    } else {
+        return IGraphTransformer::TStatus::Repeat;
+    }
+    return IGraphTransformer::TStatus::Ok;
+}
+
+} // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/type_ann.cpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <ydb/library/yql/ast/yql_expr.h>
+#include <ydb/library/yql/ast/yql_pos_handle.h>
+#include <ydb/library/yql/core/expr_nodes_gen/yql_expr_nodes_gen.h>
+#include <ydb/library/yql/core/yql_graph_transformer.h>
+
+namespace NYql::NPushdown {
+
+IGraphTransformer::TStatus AnnotateFilterPredicate(
+    const TExprNode::TPtr& input,
+    size_t childIndex,
+    const TStructExprType* itemType,
+    TExprContext& ctx);
+
+
+} // namespace NYql::NPushdown

ydb/library/yql/providers/common/pushdown/type_ann.h

@@ -4,6 +4,8 @@ SRCS(
     collection.cpp
     predicate_node.cpp
     settings.cpp
+    type_ann.cpp
+    physical_opt.cpp
 )
 
 PEERDIR(