Apply ExtractMembers to PartitionsByKeys

commit_hash:479add2c66911eda0331429ee201911cc410bcbd

Author: nepal

yql/essentials/core/common_opt/yql_co_extr_members.cpp

@@ -505,14 +505,15 @@ TExprNode::TPtr ApplyExtractMembersToFlatMap(const TExprNode::TPtr& node, const
 }
 
 TExprNode::TPtr ApplyExtractMembersToPartitionByKey(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
-    TCoPartitionByKey part(node);
+    TCoPartitionByKeyBase part(node);
     YQL_CLOG(DEBUG, Core) << "Apply ExtractMembers to " << node->Content() << logSuffix;
     auto newBody = Build<TCoExtractMembers>(ctx, part.Pos())
         .Input(part.ListHandlerLambda().Body())
         .Members(members)
         .Done();
 
-    return Build<TCoPartitionByKey>(ctx, part.Pos())
+    return Build<TCoPartitionByKeyBase>(ctx, part.Pos())
+        .CallableName(node->Content())
         .Input(part.Input())
         .KeySelectorLambda(part.KeySelectorLambda())
         .ListHandlerLambda()

yql/essentials/core/common_opt/yql_co_finalizers.cpp

@@ -566,7 +566,12 @@ void RegisterCoFinalizers(TFinalizingOptimizerMap& map) {
         return true;
     };
 
-    map[TCoPartitionByKey::CallableName()] = [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
+    map[TCoPartitionByKey::CallableName()] = map[TCoPartitionsByKeys::CallableName()] =
+        [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx)
+    {
+        if (node->IsCallable(TCoPartitionsByKeys::CallableName()) && !CanApplyExtractMembersToPartitionsByKeys(optCtx.Types)) {
+            return true;
+        }
         OptimizeSubsetFieldsForNodeWithMultiUsage(node, *optCtx.ParentsMap, toOptimize, ctx,
             [] (const TExprNode::TPtr& input, const TExprNode::TPtr& members, const TParentsMap&, TExprContext& ctx) {
                 return ApplyExtractMembersToPartitionByKey(input, members, ctx, " with multi-usage");

yql/essentials/core/common_opt/yql_co_flow2.cpp

@@ -2461,7 +2461,11 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
             return node;
         }
 
-        if (self.Input().Maybe<TCoPartitionByKey>()) {
+        if (self.Input().Maybe<TCoPartitionByKeyBase>()) {
+            if (self.Input().Maybe<TCoPartitionsByKeys>() && !CanApplyExtractMembersToPartitionsByKeys(optCtx.Types)) {
+                return node;
+            }
+
             if (auto res = ApplyExtractMembersToPartitionByKey(self.Input().Ptr(), self.Members().Ptr(), ctx, {})) {
                 return res;
             }

yql/essentials/core/yql_opt_utils.cpp

@@ -2841,4 +2841,10 @@ bool CanFuseLambdas(const TExprNode& outer, const TExprNode& inner, const TTypeA
     }
 }
 
+bool CanApplyExtractMembersToPartitionsByKeys(const TTypeAnnotationContext* types) {
+    YQL_ENSURE(types);
+    static const char optName[] = "ExtractMembersForPartitionsByKeys";
+    return IsOptimizerEnabled<optName>(*types) && !IsOptimizerDisabled<optName>(*types);
+}
+
 }

yql/essentials/core/yql_opt_utils.h

@@ -226,4 +226,6 @@ bool IsNormalizedDependsOn(const TExprNode& node);
 
 bool CanFuseLambdas(const TExprNode& outer, const TExprNode& inner, const TTypeAnnotationContext& types);
 
+bool CanApplyExtractMembersToPartitionsByKeys(const TTypeAnnotationContext* types);
+
 }

yql/essentials/tests/s-expressions/minirun/part5/canondata/result.json

@@ -352,6 +352,20 @@
             "uri": "https://{canondata_backend}/1936842/0e8399ea57053215a346dc8ce52f4206386c19fb/resource.tar.gz#test.test_InMem-CompareStruct2-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[InMem-ExtractMembersOverPartitionsByKyes-default.txt-Debug]": [
+        {
+            "checksum": "d2b4c78b235e2873373ae3d82987b842",
+            "size": 1372,
+            "uri": "https://{canondata_backend}/1600758/098090b2fdd96c8505b3c52b52768447f489d79e/resource.tar.gz#test.test_InMem-ExtractMembersOverPartitionsByKyes-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[InMem-ExtractMembersOverPartitionsByKyes-default.txt-Results]": [
+        {
+            "checksum": "486117c929f3e77513590f8dec8cbb9c",
+            "size": 837,
+            "uri": "https://{canondata_backend}/1600758/098090b2fdd96c8505b3c52b52768447f489d79e/resource.tar.gz#test.test_InMem-ExtractMembersOverPartitionsByKyes-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[InMem-Fold-default.txt-Debug]": [
         {
             "checksum": "3cf721a14fa47fd67c1967d58d82a7c2",

yql/essentials/tests/s-expressions/suites/InMem/ExtractMembersOverPartitionsByKyes.yqls

@@ -0,0 +1,36 @@
+(
+#comment
+(let config (DataSource 'config))
+(let world (Configure! world config '"OptimizerFlags" '"ExtractMembersForPartitionsByKeys"))
+
+(let res_sink (DataSink 'result))
+
+(let list (AsList
+(AsStruct '('key (Uint32 '1)) '('subkey (Uint32 '2)) '('value (String 'a)) '('unused (Int32 '1)))
+(AsStruct '('key (Uint32 '1)) '('subkey (Uint32 '1)) '('value (String 'b)) '('unused (Int32 '1)))
+(AsStruct '('key (Uint32 '2)) '('subkey (Uint32 '4)) '('value (String 'c)) '('unused (Int32 '1)))
+(AsStruct '('key (Uint32 '1)) '('subkey (Uint32 '3)) '('value (String 'd)) '('unused (Int32 '1)))
+(AsStruct '('key (Uint32 '2)) '('subkey (Uint32 '6)) '('value (String 'e)) '('unused (Int32 '1)))
+(AsStruct '('key (Uint32 '3)) '('subkey (Uint32 '5)) '('value (String 'f)) '('unused (Int32 '1)))
+))
+
+
+(let keyExtractor (lambda '(x) (Member x 'key)))
+(let sortKeyExtractor (lambda '(x) (Member x 'subkey)))
+(let direction (Bool 'true))
+
+(let handler (lambda '(x) (Condense1 x
+    (lambda '(row) (AsStruct '('key (Member row 'key)) 
+        '('group (Concat (String '"values:") (Member row 'value)))))
+    (lambda '(row state) (AggrNotEquals (Member row 'key) (Member state 'key)))
+    (lambda '(row state) (AsStruct '('key (Member state 'key)) '('group 
+    (Concat (Concat (Member state 'group) (String '" ")) (Member row 'value)))))
+)))
+(let data (PartitionsByKeys (Iterator list) keyExtractor direction sortKeyExtractor handler))
+(let data (Collect data))
+(let data (ExtractMembers data '('group)))
+
+(let world (Write! world res_sink (Key) data '('('type))))
+(let world (Commit! world res_sink))
+(return world)
+)