correct DqRewriteLengthOfStageOutputLegacy usage -> DqRewriteLengthOfStageOutput through yp precalculation implemented

canondata

6 files changed, 11 insertions, 110 deletions

diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index a3bdbedf15b..6d940cbb7cc 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -1971,108 +1971,6 @@ TExprBase DqRewriteLengthOfStageOutput(TExprBase node, TExprContext& ctx, IOptim
     return precompute;
 }
 
-// TODO: Remove once precomputes are supported in YQ
-TExprBase DqRewriteLengthOfStageOutputLegacy(TExprBase node, TExprContext& ctx, IOptimizationContext&) {
-    if (!node.Maybe<TCoLength>().List().Maybe<TDqCnUnionAll>()) {
-        return node;
-    }
-
-    auto dqUnion = node.Cast<TCoLength>().List().Cast<TDqCnUnionAll>();
-
-    auto zero = Build<TCoUint64>(ctx, node.Pos())
-                    .Literal().Build("0")
-                    .Done();
-
-    auto field = BuildAtom("_dq_agg_cnt", node.Pos(), ctx);
-
-    auto aggregateCombine = Build<TCoAggregateCombine>(ctx, node.Pos())
-        .Input(dqUnion)
-        .Keys()
-        .Build()
-        .Handlers()
-            .Add<TCoAggregateTuple>()
-                .ColumnName(field)
-                .Trait<TCoAggApply>()
-                    .Name<TCoAtom>()
-                        .Value("count_all")
-                    .Build()
-                    .InputType(ExpandType(node.Pos(), GetSeqItemType(*dqUnion.Raw()->GetTypeAnn()), ctx))
-                    .Extractor<TCoLambda>()
-                        .Args({ "row" })
-                        .Body<TCoVoid>()
-                        .Build()
-                    .Build()
-                .Build()
-            .Build()
-        .Build()
-        .Settings()
-        .Build()
-        .Done();
-
-    TVector<const TItemExprType*> stateItems = {
-        ctx.MakeType<TItemExprType>("_dq_agg_cnt", ctx.MakeType<TDataExprType>(EDataSlot::Uint64))
-    };
-
-    auto stateRowType = ctx.MakeType<TStructExprType>(stateItems);
-    auto stateTypeNode = ExpandType(node.Pos(), *stateRowType, ctx);
-
-    auto originalTypeNode = ctx.Builder(node.Pos())
-        .Callable("StructType")
-            .List(0)
-                .Add(0, field.Ptr())
-                .Callable(1, "VoidType")
-                .Seal()
-            .Seal()
-        .Seal()
-        .Build();
-
-    auto aggregateFinal = Build<TCoAggregateMergeFinalize>(ctx, node.Pos())
-        .Input(aggregateCombine)
-        .Keys()
-        .Build()
-        .Handlers()
-            .Add<TCoAggregateTuple>()
-                .ColumnName(field)
-                .Trait<TCoAggApplyState>()
-                    .Name<TCoAtom>()
-                        .Value("count_all")
-                    .Build()
-                    .InputType(stateTypeNode)
-                    .Extractor<TCoLambda>()
-                        .Args({ "row" })
-                        .Body<TCoMember>()
-                            .Struct("row")
-                            .Name(field)
-                        .Build()
-                    .Build()
-                    .OriginalType(originalTypeNode)
-                .Build()
-            .Build()
-        .Build()
-        .Settings()
-        .Build()
-        .Done();
-
-    auto toOptional = Build<TCoToOptional>(ctx, node.Pos())
-            .List(aggregateFinal)
-            .Done();
-
-    auto coalesce = Build<TCoCoalesce>(ctx, node.Pos())
-            .Predicate(toOptional)
-            .Value<TCoAsStruct>()
-                .Add<TCoNameValueTuple>()
-                    .Name(field)
-                    .Value(zero)
-                    .Build()
-                .Build()
-            .Done();
-
-    return Build<TCoMember>(ctx, node.Pos())
-        .Struct(coalesce)
-        .Name(field)
-        .Done();
-}
-
 TExprBase DqBuildPureExprStage(TExprBase node, TExprContext& ctx) {
     if (!IsDqPureExpr(node)) {
         return node;
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.h b/ydb/library/yql/dq/opt/dq_opt_phy.h
index 88be69c3c87..ca5a0bd4e9b 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.h
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.h
@@ -76,9 +76,6 @@ NNodes::TExprBase DqBuildTakeSkipStage(NNodes::TExprBase node, TExprContext& ctx
 NNodes::TExprBase DqRewriteLengthOfStageOutput(NNodes::TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage);
 
-NNodes::TExprBase DqRewriteLengthOfStageOutputLegacy(NNodes::TExprBase node, TExprContext& ctx,
-    IOptimizationContext& optCtx);
-
 NNodes::TExprBase DqRewriteRightJoinToLeft(const NNodes::TExprBase node, TExprContext& ctx);
 
 NNodes::TExprBase DqRewriteLeftPureJoin(const NNodes::TExprBase node, TExprContext& ctx,
diff --git a/ydb/library/yql/dq/type_ann/dq_type_ann.h b/ydb/library/yql/dq/type_ann/dq_type_ann.h
index afcc29eb88e..9c96d67b8e1 100644
--- a/ydb/library/yql/dq/type_ann/dq_type_ann.h
+++ b/ydb/library/yql/dq/type_ann/dq_type_ann.h
@@ -8,7 +8,8 @@
 namespace NYql::NDq {
 
 IGraphTransformer::TStatus AnnotateDqStage(const TExprNode::TPtr& input, TExprContext& ctx);
-IGraphTransformer::TStatus AnnotateDqPhyStage(const TExprNode::TPtr& input, TExprContext& ctx);
+IGraphTransformer::TStatus AnnotateDqPhyLength(const TExprNode::TPtr& node, TExprContext& ctx);
+IGraphTransformer::TStatus AnnotateDqPhyStage(const TExprNode::TPtr & input, TExprContext & ctx);
 IGraphTransformer::TStatus AnnotateDqOutput(const TExprNode::TPtr& input, TExprContext& ctx);
 IGraphTransformer::TStatus AnnotateDqCnHashShuffle(const TExprNode::TPtr& input, TExprContext& ctx);
 IGraphTransformer::TStatus AnnotateDqCnValue(const TExprNode::TPtr& input, TExprContext& ctx);
diff --git a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp
index 69debeb3f13..cbb44b135f8 100644
--- a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp
+++ b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp
@@ -41,7 +41,7 @@ public:
         AddHandler(0, &TCoTopSort::Match, HNDL(BuildTopSortStage<false>));
         AddHandler(0, &TCoSort::Match, HNDL(BuildSortStage<false>));
         AddHandler(0, &TCoTakeBase::Match, HNDL(BuildTakeOrTakeSkipStage<false>));
-        AddHandler(0, &TCoLength::Match, HNDL(RewriteLengthOfStageOutput));
+        AddHandler(0, &TCoLength::Match, HNDL(RewriteLengthOfStageOutput<false>));
         AddHandler(0, &TCoExtendBase::Match, HNDL(BuildExtendStage));
         AddHandler(0, &TDqJoin::Match, HNDL(RewriteRightJoinToLeft));
         AddHandler(0, &TDqJoin::Match, HNDL(RewriteLeftPureJoin<false>));
@@ -73,6 +73,7 @@ public:
         AddHandler(1, &TCoTopSort::Match, HNDL(BuildTopSortStage<true>));
         AddHandler(1, &TCoSort::Match, HNDL(BuildSortStage<true>));
         AddHandler(1, &TCoTakeBase::Match, HNDL(BuildTakeOrTakeSkipStage<true>));
+        AddHandler(1, &TCoLength::Match, HNDL(RewriteLengthOfStageOutput<true>));
         AddHandler(1, &TDqJoin::Match, HNDL(RewriteLeftPureJoin<true>));
         AddHandler(1, &TDqJoin::Match, HNDL(BuildJoin<true>));
         AddHandler(1, &TCoAssumeSorted::Match, HNDL(BuildSortStage<true>));
@@ -318,8 +319,9 @@ protected:
         }
     }
 
-    TMaybeNode<TExprBase> RewriteLengthOfStageOutput(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx) {
-        return DqRewriteLengthOfStageOutputLegacy(node, ctx, optCtx);
+    template <bool IsGlobal>
+    TMaybeNode<TExprBase> RewriteLengthOfStageOutput(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx, const TGetParents& getParents) {
+        return DqRewriteLengthOfStageOutput(node, ctx, optCtx, *getParents(), IsGlobal);
     }
 
     TMaybeNode<TExprBase> BuildExtendStage(TExprBase node, TExprContext& ctx) {
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_datasource_constraints.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_datasource_constraints.cpp
index 914c9156259..52399454561 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_datasource_constraints.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_datasource_constraints.cpp
@@ -23,7 +23,8 @@ public:
             TDqSource::CallableName(),
             TDqSourceWrap::CallableName(),
             TDqSourceWideWrap::CallableName(),
-            TDqSourceWideBlockWrap::CallableName()
+            TDqSourceWideBlockWrap::CallableName(),
+            TDqPhyLength::CallableName()
         }, Hndl(&TDqDataSourceConstraintTransformer::HandleDefault));
     }
 
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_datasource_type_ann.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_datasource_type_ann.cpp
index ee54ec50743..60ff588c0f7 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_datasource_type_ann.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_datasource_type_ann.cpp
@@ -25,6 +25,8 @@ public:
         AddHandler({TDqReadWideWrap::CallableName()}, Hndl(&TDqsDataSourceTypeAnnotationTransformer::HandleWideReadWrap));
         AddHandler({TCoConfigure::CallableName()}, Hndl(&TDqsDataSourceTypeAnnotationTransformer::HandleConfig));
         AddHandler({TDqSource::CallableName()}, Hndl(&NDq::AnnotateDqSource));
+        AddHandler({TDqPhyLength::CallableName()}, Hndl(&NDq::AnnotateDqPhyLength));
+        
     }
 
 private: