YQ-3617: fix GROUP BY HOP + AS_TABLE (ydb-platform#9370)

ydb/library/yql/core/common_opt/yql_co_simple1.cpp

@@ -5,6 +5,7 @@
 #include <ydb/library/yql/core/yql_atom_enums.h>
 #include <ydb/library/yql/core/yql_expr_type_annotation.h>
 #include <ydb/library/yql/core/yql_join.h>
+#include <ydb/library/yql/core/yql_opt_hopping.h>
 #include <ydb/library/yql/core/yql_opt_utils.h>
 #include <ydb/library/yql/core/yql_opt_window.h>
 #include <ydb/library/yql/core/yql_type_helpers.h>
@@ -3043,6 +3044,190 @@ TExprNode::TPtr Normalize(const TCoAggregate& node, TExprContext& ctx) {
         .Done().Ptr();
 }
 
+TExprNode::TPtr RemoveDeadPayloadColumns(const TCoAggregate& aggr, TExprContext& ctx) {
+    const TExprNode::TPtr outputColumnsSetting = GetSetting(aggr.Settings().Ref(), "output_columns");
+    if (!outputColumnsSetting) {
+        return aggr.Ptr();
+    }
+    const TExprNode::TPtr outputColumns = outputColumnsSetting->ChildPtr(1);
+
+    TSet<TStringBuf> outMembers;
+    for (const auto& x : outputColumns->ChildrenList()) {
+        outMembers.insert(x->Content());
+    }
+
+    TExprNodeList newHandlers;
+    bool rebuildHandlers = false;
+    for (const auto& handler : aggr.Handlers()) {
+        if (handler.ColumnName().Ref().IsList()) {
+            // many columns
+            auto columns = handler.ColumnName().Cast<TCoAtomList>();
+            TVector<size_t> liveIndexes;
+            for (size_t i = 0; i < columns.Size(); ++i) {
+                if (outMembers.contains(columns.Item(i).Value())) {
+                    liveIndexes.push_back(i);
+                }
+            }
+
+            if (liveIndexes.empty()) {
+                // drop handler
+                rebuildHandlers = true;
+                continue;
+            } else if (liveIndexes.size() < columns.Size() && handler.Trait().Maybe<TCoAggregationTraits>()) {
+                auto traits = handler.Trait().Cast<TCoAggregationTraits>();
+
+                TExprNodeList nameNodes;
+                TExprNodeList finishBody;
+                TExprNode::TPtr arg = ctx.NewArgument(traits.FinishHandler().Pos(), "arg");
+                auto originalTuple = ctx.Builder(traits.FinishHandler().Pos())
+                    .Apply(traits.FinishHandler().Ref())
+                        .With(0, arg)
+                    .Seal()
+                    .Build();
+
+                for (auto& idx : liveIndexes) {
+                    nameNodes.emplace_back(columns.Item(idx).Ptr());
+                    finishBody.emplace_back(ctx.Builder(traits.FinishHandler().Pos())
+                        .Callable("Nth")
+                            .Add(0, originalTuple)
+                            .Atom(1, idx)
+                        .Seal()
+                        .Build());
+                }
+
+                auto finishLambda = ctx.NewLambda(traits.FinishHandler().Pos(),
+                    ctx.NewArguments(traits.FinishHandler().Pos(), { arg }),
+                    ctx.NewList(traits.FinishHandler().Pos(), std::move(finishBody)));
+
+                auto newHandler = Build<TCoAggregateTuple>(ctx, handler.Pos())
+                    .InitFrom(handler)
+                    .ColumnName(ctx.NewList(handler.ColumnName().Pos(), std::move(nameNodes)))
+                    .Trait<TCoAggregationTraits>()
+                        .InitFrom(traits)
+                        .FinishHandler(finishLambda)
+                    .Build()
+                    .Done().Ptr();
+
+                newHandlers.emplace_back(std::move(newHandler));
+                rebuildHandlers = true;
+                continue;
+            }
+        } else {
+            if (!outMembers.contains(handler.ColumnName().Cast<TCoAtom>().Value())) {
+                // drop handler
+                rebuildHandlers = true;
+                continue;
+            }
+        }
+
+        newHandlers.push_back(handler.Ptr());
+    }
+
+    if (rebuildHandlers) {
+        YQL_CLOG(DEBUG, Core) << "Drop unused payloads in " << aggr.CallableName();
+        return Build<TCoAggregate>(ctx, aggr.Pos())
+            .InitFrom(aggr)
+            .Handlers(ctx.NewList(aggr.Pos(), std::move(newHandlers)))
+            .Done()
+            .Ptr();
+    }
+
+    return aggr.Ptr();
+}
+
+TExprNode::TPtr RewriteAsHoppingWindowFullOutput(const TCoAggregate& aggregate, TExprContext& ctx) {
+    const auto pos = aggregate.Pos();
+
+    NHopping::EnsureNotDistinct(aggregate);
+
+    const auto maybeHopTraits = NHopping::ExtractHopTraits(aggregate, ctx, false);
+    if (!maybeHopTraits) {
+        return nullptr;
+    }
+    const auto hopTraits = *maybeHopTraits;
+
+    const auto aggregateInputType = GetSeqItemType(*aggregate.Ptr()->Head().GetTypeAnn()).Cast<TStructExprType>();
+    NHopping::TKeysDescription keysDescription(*aggregateInputType, aggregate.Keys(), hopTraits.Column);
+
+    const auto keyLambda = keysDescription.GetKeySelector(ctx, pos, aggregateInputType);
+    const auto timeExtractorLambda = NHopping::BuildTimeExtractor(hopTraits.Traits, ctx);
+    const auto initLambda = NHopping::BuildInitHopLambda(aggregate, ctx);
+    const auto updateLambda = NHopping::BuildUpdateHopLambda(aggregate, ctx);
+    const auto saveLambda = NHopping::BuildSaveHopLambda(aggregate, ctx);
+    const auto loadLambda = NHopping::BuildLoadHopLambda(aggregate, ctx);
+    const auto mergeLambda = NHopping::BuildMergeHopLambda(aggregate, ctx);
+    const auto finishLambda = NHopping::BuildFinishHopLambda(aggregate, keysDescription.GetActualGroupKeys(), hopTraits.Column, ctx);
+
+    const auto streamArg = Build<TCoArgument>(ctx, pos).Name("stream").Done();
+    auto multiHoppingCoreBuilder = Build<TCoMultiHoppingCore>(ctx, pos)
+        .KeyExtractor(keyLambda)
+        .TimeExtractor(timeExtractorLambda)
+        .Hop(hopTraits.Traits.Hop())
+        .Interval(hopTraits.Traits.Interval())
+        .Delay(hopTraits.Traits.Delay())
+        .DataWatermarks(hopTraits.Traits.DataWatermarks())
+        .InitHandler(initLambda)
+        .UpdateHandler(updateLambda)
+        .MergeHandler(mergeLambda)
+        .FinishHandler(finishLambda)
+        .SaveHandler(saveLambda)
+        .LoadHandler(loadLambda)
+        .template WatermarkMode<TCoAtom>().Build(ToString(false));
+
+    return Build<TCoPartitionsByKeys>(ctx, pos)
+        .Input(aggregate.Input())
+        .KeySelectorLambda(keyLambda)
+        .SortDirections<TCoBool>()
+                .Literal()
+                .Value("true")
+                .Build()
+            .Build()
+        .SortKeySelectorLambda(timeExtractorLambda)
+        .ListHandlerLambda()
+            .Args(streamArg)
+            .template Body<TCoForwardList>()
+                .Stream(Build<TCoMap>(ctx, pos)
+                    .Input(multiHoppingCoreBuilder
+                        .template Input<TCoIterator>()
+                            .List(streamArg)
+                            .Build()
+                        .Done())
+                    .Lambda(keysDescription.BuildUnpickleLambda(ctx, pos, *aggregateInputType))
+                    .Done())
+                .Build()
+            .Build()
+        .Done()
+        .Ptr();
+}
+
+TExprNode::TPtr RewriteAsHoppingWindow(TExprNode::TPtr node, TExprContext& ctx) {
+    const auto aggregate = TCoAggregate(node);
+
+    if (!IsPureIsolatedLambda(*aggregate.Ptr())) {
+        return nullptr;
+    }
+
+    if (!GetSetting(aggregate.Settings().Ref(), "hopping")) {
+        return nullptr;
+    }
+
+    auto result = RewriteAsHoppingWindowFullOutput(aggregate, ctx);
+    if (!result) {
+        return result;
+    }
+
+    auto outputColumnSetting = GetSetting(aggregate.Settings().Ref(), "output_columns");
+    if (!outputColumnSetting) {
+        return result;
+    }
+
+    return Build<TCoExtractMembers>(ctx, aggregate.Pos())
+        .Input(result)
+        .Members(outputColumnSetting->ChildPtr(1))
+        .Done()
+        .Ptr();
+}
+
 TExprNode::TPtr PullAssumeColumnOrderOverEquiJoin(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
     TVector<ui32> withAssume;
     for (ui32 i = 0; i < node->ChildrenSize() - 2; i++) {
@@ -4763,6 +4948,15 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
             return normalized;
         }
 
+        if (auto clean = RemoveDeadPayloadColumns(self, ctx); clean != node) {
+            return clean;
+        }
+
+        if (auto hopping = RewriteAsHoppingWindow(node, ctx)) {
+            YQL_CLOG(DEBUG, Core) << "RewriteAsHoppingWindow";
+            return hopping;
+        }
+
         return DropReorder<false>(node, ctx);
     };
 

ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp

@@ -3418,6 +3418,35 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TUniqueConstraintNode>(exprRoot, "LazyList", "");
         CheckConstraint<TDistinctConstraintNode>(exprRoot, "LazyList", "");
     }
+
+    Y_UNIT_TEST(GroupByHop) {
+        const TStringBuf s = R"((
+(let list (AsList
+    (AsStruct '('"time" (String '"2024-01-01T00:00:01Z")) '('"user" (Int32 '"1")) '('"data" (Null)))
+    (AsStruct '('"time" (String '"2024-01-01T00:00:02Z")) '('"user" (Int32 '"1")) '('"data" (Null)))
+    (AsStruct '('"time" (String '"2024-01-01T00:00:03Z")) '('"user" (Int32 '"1")) '('"data" (Null)))
+))
+(let input (FlatMap list (lambda '(row) (Just (AsStruct '('"data" (Member row '"data")) '('group0 (AsList (Member row '"user"))) '('"time" (Member row '"time")) '('"user" (Member row '"user")))))))
+(let keySelector (lambda '(row) '((StablePickle (Member row '"data")) (StablePickle (Member row 'group0)))))
+(let sortKeySelector (lambda '(row) (SafeCast (Member row '"time") (OptionalType (DataType 'Timestamp)))))
+(let res (PartitionsByKeys input keySelector (Bool 'true) sortKeySelector (lambda '(row) (block '(
+  (let interval (Interval '1000000))
+  (let map (lambda '(item) (AsStruct)))
+  (let reduce (lambda '(lhs rhs) (AsStruct)))
+  (let hopping (MultiHoppingCore (Iterator row) keySelector sortKeySelector interval interval interval 'true map reduce map map reduce (lambda '(key state time) (AsStruct '('_yql_time time) '('"data" (Nth key '"0")) '('group0 (Nth key '"1")))) '"0"))
+  (return (ForwardList (FlatMap hopping (lambda '(row) (Just (AsStruct '('_yql_time (Member row '_yql_time)) '('"data" (Unpickle (NullType) (Member row '"data"))) '('group0 (Unpickle (ListType (DataType 'Int32)) (Member row 'group0)))))))))
+)))))
+
+(let res_sink (DataSink 'yt (quote plato)))
+(let world (Write! world res_sink (Key '('table (String 'Output))) res '('('mode 'renew))))
+(return (Commit! world res_sink))
+        ))";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TDistinctConstraintNode>(exprRoot, "PartitionsByKeys", "Distinct((data,group0))");
+        CheckConstraint<TUniqueConstraintNode>(exprRoot, "PartitionsByKeys", "Unique((data,group0))");
+    }
 }
 
 } // namespace NYql

ydb/library/yql/core/ya.make

@@ -28,6 +28,7 @@ SRCS(
     yql_join.cpp
     yql_join.h
     yql_library_compiler.cpp
+    yql_opt_hopping.cpp
     yql_opt_match_recognize.cpp
     yql_opt_match_recognize.h
     yql_opt_proposed_by_data.cpp

ydb/library/yql/core/yql_expr_constraint.cpp

@@ -239,6 +239,9 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
         Functions["BlockExpandChunked"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
         Functions["BlockMergeFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
         Functions["BlockMergeManyFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
+        Functions["MultiHoppingCore"] = &TCallableConstraintTransformer::MultiHoppingCoreWrap;
+        Functions["StablePickle"] = &TCallableConstraintTransformer::FromFirst<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
+        Functions["Unpickle"] = &TCallableConstraintTransformer::FromSecond<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
     }
 
     std::optional<IGraphTransformer::TStatus> ProcessCore(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
@@ -3144,6 +3147,26 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
 
         return TStatus::Ok;
     }
+
+    TStatus MultiHoppingCoreWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
+        if (const auto status = UpdateAllChildLambdasConstraints(*input); status != TStatus::Ok) {
+            return status;
+        }
+
+        TExprNode::TPtr keySelectorLambda = input->Child(TCoMultiHoppingCore::idx_KeyExtractor);
+        const auto keys = GetPathsToKeys(keySelectorLambda->Tail(), keySelectorLambda->Head().Head());
+        std::vector<std::string_view> columns(keys.size());
+        std::transform(keys.begin(), keys.end(), columns.begin(), [](const TPartOfConstraintBase::TPathType& path) -> std::string_view {
+            return path.front();
+        });
+        if (!columns.empty()) {
+            input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(columns));
+            input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(columns));
+        }
+
+        return TStatus::Ok;
+    }
+
 private:
     template <class TConstraintContainer>
     static void CopyExcept(TConstraintContainer& dst, const TConstraintContainer& from, const TSet<TStringBuf>& except) {