[dq] Add Top optimizer (ydb-platform#10104)

GrigoriyPA · Oct 4, 2024 · 01ce21a · 01ce21a
1 parent 2e6d05a
commit 01ce21a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp
@@ -38,6 +38,7 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         AddHandler(0, &TCoPartitionByKey::Match, HNDL(BuildPartitionStage<false>));
         AddHandler(0, &TCoAsList::Match, HNDL(BuildAggregationResultStage));
         AddHandler(0, &TCoTopSort::Match, HNDL(BuildTopSortStage<false>));
+        AddHandler(0, &TCoTop::Match, HNDL(BuildTopStage<false>));
         AddHandler(0, &TCoSort::Match, HNDL(BuildSortStage<false>));
         AddHandler(0, &TCoTakeBase::Match, HNDL(BuildTakeOrTakeSkipStage<false>));
         AddHandler(0, &TCoLength::Match, HNDL(RewriteLengthOfStageOutput<false>));
@@ -74,6 +75,7 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         AddHandler(1, &TCoFinalizeByKey::Match, HNDL(BuildFinalizeByKeyStage<true>));
         AddHandler(1, &TCoPartitionByKey::Match, HNDL(BuildPartitionStage<true>));
         AddHandler(1, &TCoTopSort::Match, HNDL(BuildTopSortStage<true>));
+        AddHandler(1, &TCoTop::Match, HNDL(BuildTopStage<true>));
         AddHandler(1, &TCoSort::Match, HNDL(BuildSortStage<true>));
         AddHandler(1, &TCoTakeBase::Match, HNDL(BuildTakeOrTakeSkipStage<true>));
         AddHandler(1, &TCoLength::Match, HNDL(RewriteLengthOfStageOutput<true>));
@@ -184,6 +186,11 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         return DqBuildTopSortStage(node, ctx, optCtx, *getParents(), IsGlobal);
     }
 
+    template <bool IsGlobal>
+    TMaybeNode<TExprBase> BuildTopStage(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx, const TGetParents& getParents) {
+        return DqBuildTopStage(node, ctx, optCtx, *getParents(), IsGlobal);
+    }
+
     template <bool IsGlobal>
     TMaybeNode<TExprBase> BuildSortStage(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx, const TGetParents& getParents) {
         return DqBuildSortStage(node, ctx, optCtx, *getParents(), IsGlobal);