Prohibit UDF on Flow in case of PartitionByKeys callable. (#6322)

Author: Darych

ydb/library/yql/core/yql_opt_utils.cpp

@@ -1813,7 +1813,8 @@ TExprNode::TPtr FindNonYieldTransparentNodeImpl(const TExprNode::TPtr& root, con
                 || TCoForwardList::Match(node.Get())
                 || TCoApply::Match(node.Get())
                 || TCoSwitch::Match(node.Get())
-                || node->IsCallable("DqReplicate");
+                || node->IsCallable("DqReplicate")
+                || TCoPartitionsByKeys::Match(node.Get());
         }
     );
 
@@ -1851,6 +1852,11 @@ TExprNode::TPtr FindNonYieldTransparentNodeImpl(const TExprNode::TPtr& root, con
                     return node;
                 }
             }
+        } else if (TCoPartitionsByKeys::Match(candidate.Get())) {
+            const auto handlerChild = candidate->Child(TCoPartitionsByKeys::idx_ListHandlerLambda);
+            if (auto node = FindNonYieldTransparentNodeImpl(handlerChild->TailPtr(), udfSupportsYield, TNodeSet{&handlerChild->Head().Head()})) {
+                return node;
+            }
         }
     }
     return {};

ydb/library/yql/tests/sql/suites/produce/reduce_all_with_python_input_stream-dq_fail.cfg

@@ -0,0 +1,4 @@
+xfail
+in Input1 input1.txt
+udf python3_udf
+providers dq

ydb/library/yql/tests/sql/suites/produce/reduce_all_with_python_input_stream._sql

@@ -0,0 +1,14 @@
+/* postgres can not */
+USE plato;
+
+$udfScript = @@
+import functools
+def Len(stream):
+    sums = [functools.reduce(lambda x,y: x + y, pair[1], 0) for pair in stream]
+    return {"sumByAllVal":functools.reduce(lambda x,y: x + y, sums, 0)}
+@@;
+
+$udf = Python3::Len(Callable<(Stream<Tuple<String,Stream<Uint32>>>)->Struct<sumByAllVal:Uint32>>, $udfScript);
+
+--INSERT INTO Output
+REDUCE Input1 ON key USING ALL $udf(cast(value as uint32) ?? 0);

ydb/library/yql/tests/sql/suites/produce/reduce_all_with_python_input_stream.cfg

@@ -0,0 +1,3 @@
+in Input1 input1.txt
+udf python3_udf
+providers yt

ydb/library/yql/tests/sql/suites/produce/reduce_with_python_input_stream._sql

@@ -0,0 +1,14 @@
+/* postgres can not */
+USE plato;
+
+$udfScript = @@
+import functools
+def Len(key, input):
+    return {"value":functools.reduce(lambda x,y: x + 1, input, 0)}
+@@;
+
+$udf = Python::Len(Callable<(String, Stream<String>)->Struct<value:Uint32>>, $udfScript);
+
+$res = (REDUCE Input1 ON key USING $udf(value));
+
+select * from $res order by value;

ydb/library/yql/tests/sql/suites/produce/reduce_with_python_input_stream.cfg

@@ -0,0 +1,2 @@
+in Input1 input1.txt
+udf python3_udf