AggregateSubsetFieldsAnalyzer

1 files changed, 39 insertions, 2 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
index b4c1b618b9e..a606a61b887 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp
@@ -50,11 +50,16 @@ TExprNode::TPtr AggregateSubsetFieldsAnalyzer(const TCoAggregate& node, TExprCon
         }
         else {
             auto traits = x.Ref().Child(1);
-            if (!traits->IsCallable("AggregationTraits")) {
+            ui32 index;
+            if (traits->IsCallable("AggregationTraits")) {
+                index = 0;
+            } else if (traits->IsCallable("AggApply")) {
+                index = 1;
+            } else {
                 return node.Ptr();
             }
 
-            auto structType = traits->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
+            auto structType = traits->Child(index)->GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
             for (const auto& item : structType->GetItems()) {
                 usedFields.insert(item->GetName());
             }
@@ -2398,6 +2403,38 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
             .Build();
     };
 
+    map["AggApply"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        auto type = node->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+        if (type->GetKind() != ETypeAnnotationKind::Struct) {
+            // usually distinct, type of column is used instead
+            return node;
+        }
+
+        auto structType = type->Cast<TStructExprType>();
+        TSet<TStringBuf> usedFields;
+        auto extractor = node->Child(2);
+        TSet<TStringBuf> lambdaSubset;
+        if (!HaveFieldsSubset(extractor->ChildPtr(1), *extractor->Child(0)->Child(0), lambdaSubset, *optCtx.ParentsMap)) {
+            return node;
+        }
+
+        usedFields.insert(lambdaSubset.cbegin(), lambdaSubset.cend());
+        if (usedFields.size() == structType->GetSize()) {
+            return node;
+        }
+
+        TVector<const TItemExprType*> subsetItems;
+        for (const auto& item : structType->GetItems()) {
+            if (usedFields.contains(item->GetName())) {
+                subsetItems.push_back(item);
+            }
+        }
+
+        auto subsetType = ctx.MakeType<TStructExprType>(subsetItems);
+        YQL_CLOG(DEBUG, Core) << "FieldSubset for AggApply";
+        return ctx.ChangeChild(*node, 1, ExpandType(node->Pos(), *subsetType, ctx));
+    };
+
     map["SessionWindowTraits"] = map["SortTraits"] = map["Lag"] = map["Lead"] = map["RowNumber"] = map["Rank"] = map["DenseRank"] =
         [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx)
     {