[YQL-16903] Support BlockEngine=auto/force/disable (peephole + channels)

ydb/library/yql/core/expr_nodes/yql_expr_nodes.json

@@ -2443,6 +2443,11 @@
             "Base": "TCoInputBase",
             "Match": {"Type": "Callable", "Name": "WideFromBlocks"}
         },
+        {
+            "Name": "TCoWideToBlocks",
+            "Base": "TCoInputBase",
+            "Match": {"Type": "Callable", "Name": "WideToBlocks"}
+        },
         {
             "Name": "TCoPgSelect",
             "Base": "TCallable",

ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp

@@ -5488,6 +5488,19 @@ bool CollectBlockRewrites(const TMultiExprType* multiInputType, bool keepInputCo
     return true;
 }
 
+bool CanRewriteToBlocksWithInput(const TExprNode& input, const TTypeAnnotationContext& types) {
+    EBlockEngineMode effectiveMode = types.UseBlocks ? EBlockEngineMode::Force : types.BlockEngineMode;
+    switch (effectiveMode) {
+        case NYql::EBlockEngineMode::Disable:
+            return false;
+        case NYql::EBlockEngineMode::Auto:
+            return input.IsCallable("WideFromBlocks");
+        case NYql::EBlockEngineMode::Force:
+            return true;
+    }
+    Y_UNREACHABLE();
+}
+
 TExprNode::TPtr OptimizeWideMapBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
     const auto lambda = node->TailPtr();
     if (node->Head().IsCallable("WideFromBlocks")) {
@@ -5504,6 +5517,10 @@ TExprNode::TPtr OptimizeWideMapBlocks(const TExprNode::TPtr& node, TExprContext&
         }
     }
 
+    if (!CanRewriteToBlocksWithInput(node->Head(), types)) {
+        return node;
+    }
+
     auto multiInputType = node->Head().GetTypeAnn()->Cast<TFlowExprType>()->GetItemType()->Cast<TMultiExprType>();
     ui32 newNodes;
     TNodeMap<size_t> rewritePositions;
@@ -5541,6 +5558,10 @@ TExprNode::TPtr OptimizeWideMapBlocks(const TExprNode::TPtr& node, TExprContext&
 }
 
 TExprNode::TPtr OptimizeWideFilterBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
+    if (!CanRewriteToBlocksWithInput(node->Head(), types)) {
+        return node;
+    }
+
     auto multiInputType = node->Head().GetTypeAnn()->Cast<TFlowExprType>()->GetItemType()->Cast<TMultiExprType>();
     auto lambda = node->ChildPtr(1);
     YQL_ENSURE(lambda->ChildrenSize() == 2); // filter lambda should have single output
@@ -5661,6 +5682,10 @@ TExprNode::TPtr OptimizeSkipTakeToBlocks(const TExprNode::TPtr& node, TExprConte
         return node;
     }
 
+    if (!CanRewriteToBlocksWithInput(node->Head(), types)) {
+        return node;
+    }
+
     TStringBuf newName = node->Content() == "Skip" ? "WideSkipBlocks" : "WideTakeBlocks";
     YQL_CLOG(DEBUG, CorePeepHole) << "Convert " << node->Content() << " to " << newName;
     return ctx.Builder(node->Pos())
@@ -5701,6 +5726,10 @@ TExprNode::TPtr OptimizeTopOrSortBlocks(const TExprNode::TPtr& node, TExprContex
         return node;
     }
 
+    if (!CanRewriteToBlocksWithInput(node->Head(), types)) {
+        return node;
+    }
+
     TString newName = node->Content() + TString("Blocks");
     YQL_CLOG(DEBUG, CorePeepHole) << "Convert " << node->Content() << " to " << newName;
     auto children = node->ChildrenList();
@@ -7445,6 +7474,14 @@ TExprNode::TPtr DropToFlowDeps(const TExprNode::TPtr& node, TExprContext& ctx) {
     return ctx.ChangeChildren(*node, std::move(children));
 }
 
+TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext&) {
+    if (node->ChildrenSize() == 1 && node->Head().IsCallable("FromFlow")) {
+        YQL_CLOG(DEBUG, CorePeepHole) << "Drop ToFlow over FromFlow";
+        return node->Head().HeadPtr();
+    }
+    return node;
+}
+
 TExprNode::TPtr BuildCheckedBinaryOpOverDecimal(TPositionHandle pos, TStringBuf op, const TExprNode::TPtr& lhs, const TExprNode::TPtr& rhs, const TTypeAnnotationNode& resultType, TExprContext& ctx) {
     auto typeNode = ExpandType(pos, resultType, ctx);
     return ctx.Builder(pos)
@@ -7701,6 +7738,7 @@ struct TPeepHoleRules {
         {"AggrGreater", &ExpandAggrCompare<false, false>},
         {"AggrLessOrEqual", &ExpandAggrCompare<true, true>},
         {"AggrGreaterOrEqual", &ExpandAggrCompare<false, true>},
+        {"ToFlow", &OptimizeToFlow},
     };
 
     const TExtPeepHoleOptimizerMap FinalStageExtRules = {};