1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
#pragma once
#include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
#include <yql/essentials/core/yql_opt_utils.h>
#include "yql_type_annotation.h"
namespace NYql {
class TAggregateExpander {
public:
TAggregateExpander(bool usePartitionsByKeys, const bool useFinalizeByKeys, const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
bool forceCompact = false, bool compactForDistinct = false, bool usePhases = false, bool useBlocks = false)
: Node(node)
, Ctx(ctx)
, TypesCtx(typesCtx)
, UsePartitionsByKeys(usePartitionsByKeys)
, UseFinalizeByKeys(useFinalizeByKeys)
, ForceCompact(forceCompact)
, CompactForDistinct(compactForDistinct)
, UsePhases(usePhases)
, AggregatedColumns(nullptr)
, VoidNode(ctx.NewCallable(node->Pos(), "Void", {}))
, HaveDistinct(false)
, EffectiveCompact(false)
, HaveSessionSetting(false)
, OriginalRowType(nullptr)
, RowType(nullptr)
, UseBlocks(useBlocks)
{
PreMap = Ctx.Builder(node->Pos())
.Lambda()
.Param("premap")
.Callable("Just").Arg(0, "premap").Seal()
.Seal().Build();
SortParams = {
.Key = VoidNode,
.Order = VoidNode
};
}
TExprNode::TPtr ExpandAggregate();
static TExprNode::TPtr CountAggregateRewrite(const NNodes::TCoAggregate& node, TExprContext& ctx, bool useBlocks);
private:
using TIdxSet = std::set<ui32>;
TExprNode::TPtr ExpandAggregateWithFullOutput();
TExprNode::TPtr ExpandAggApply(const TExprNode::TPtr& node);
bool CollectTraits();
TExprNode::TPtr RebuildAggregate();
TExprNode::TPtr GetContextLambda();
void ProcessSessionSetting(TExprNode::TPtr sessionSetting);
TVector<const TTypeAnnotationNode*> GetKeyItemTypes();
bool IsNeedPickle(const TVector<const TTypeAnnotationNode*>& keyItemTypes);
TExprNode::TPtr GetKeyExtractor(bool needPickle);
void CollectColumnsSpecs();
void BuildNothingStates();
// Partial aggregate generation
TExprNode::TPtr GeneratePartialAggregate(const TExprNode::TPtr keyExtractor, const TVector<const TTypeAnnotationNode*>& keyItemTypes, bool needPickle);
TExprNode::TPtr GeneratePartialAggregateForNonDistinct(const TExprNode::TPtr& keyExtractor, const TExprNode::TPtr& pickleTypeNode);
TExprNode::TPtr GenerateDistinctGrouper(const TExprNode::TPtr distinctField,
const TVector<const TTypeAnnotationNode*>& keyItemTypes, bool needDistinctPickle);
TExprNode::TPtr ReturnKeyAsIsForCombineInit(const TExprNode::TPtr& pickleTypeNode);
// Post aggregate
TExprNode::TPtr GeneratePostAggregate(const TExprNode::TPtr& preAgg, const TExprNode::TPtr& keyExtractor);
TExprNode::TPtr GeneratePreprocessLambda(const TExprNode::TPtr& keyExtractor);
TExprNode::TPtr GenerateCondenseSwitch(const TExprNode::TPtr& keyExtractor);
TExprNode::TPtr BuildFinalizeByKeyLambda(const TExprNode::TPtr& preprocessLambda, const TExprNode::TPtr& keyExtractor);
TExprNode::TPtr GeneratePostAggregateInitPhase();
TExprNode::TPtr GeneratePostAggregateSavePhase();
TExprNode::TPtr GeneratePostAggregateMergePhase();
std::function<TExprNodeBuilder& (TExprNodeBuilder&)> GetPartialAggArgExtractor(ui32 i, bool deserialize);
TExprNode::TPtr GetFinalAggStateExtractor(ui32 i);
TExprNode::TPtr GeneratePhases();
void GenerateInitForDistinct(TExprNodeBuilder& parent, ui32& ndx, const TIdxSet& indicies, const TExprNode::TPtr& distinctField);
TExprNode::TPtr GenerateJustOverStates(const TExprNode::TPtr& input, const TIdxSet& indicies);
TExprNode::TPtr SerializeIdxSet(const TIdxSet& indicies);
TExprNode::TPtr TryGenerateBlockCombineAllOrHashed();
TExprNode::TPtr TryGenerateBlockMergeFinalizeHashed();
TExprNode::TPtr TryGenerateBlockCombine();
TExprNode::TPtr TryGenerateBlockMergeFinalize();
TExprNode::TPtr MakeInputBlocks(const TExprNode::TPtr& stream, TExprNode::TListType& keyIdxs,
TVector<TString>& outputColumns, TExprNode::TListType& aggs, bool overState, bool many, ui32* streamIdxColumn = nullptr);
private:
static constexpr TStringBuf SessionStartMemberName = "_yql_group_session_start";
const TExprNode::TPtr Node;
TExprContext& Ctx;
TTypeAnnotationContext& TypesCtx;
bool UsePartitionsByKeys;
bool UseFinalizeByKeys = false;
bool ForceCompact;
bool CompactForDistinct;
bool UsePhases;
TStringBuf Suffix;
TSessionWindowParams SessionWindowParams;
TExprNode::TPtr AggList;
TExprNode::TListType Traits;
TExprNode::TPtr KeyColumns;
TExprNode::TPtr AggregatedColumns;
const TExprNode::TPtr VoidNode;
TMaybe<TStringBuf> SessionOutputColumn;
TSortParams SortParams;
bool HaveDistinct;
bool EffectiveCompact;
bool HaveSessionSetting;
const TStructExprType* OriginalRowType;
const TStructExprType* RowType;
TVector<const TItemExprType*> RowItems;
TExprNode::TPtr PreMap;
bool UseBlocks;
TExprNode::TListType InitialColumnNames;
TExprNode::TListType FinalColumnNames;
TExprNode::TListType DistinctFields;
TExprNode::TListType NothingStates;
std::unordered_map<std::string_view, TIdxSet> Distinct2Columns;
TIdxSet NonDistinctColumns;
std::unordered_map<std::string_view, bool> DistinctFieldNeedsPickle;
std::unordered_map<std::string_view, TExprNode::TPtr> UdfSetCreate;
std::unordered_map<std::string_view, TExprNode::TPtr> UdfAddValue;
std::unordered_map<std::string_view, TExprNode::TPtr> UdfWasChanged;
};
inline TExprNode::TPtr ExpandAggregatePeepholeImpl(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
const bool useFinalizeByKey, const bool useBlocks, const bool allowSpilling) {
TAggregateExpander aggExpander(!useFinalizeByKey && !useBlocks, useFinalizeByKey, node, ctx, typesCtx,
true, false, false, typesCtx.IsBlockEngineEnabled() && !allowSpilling);
return aggExpander.ExpandAggregate();
}
TExprNode::TPtr ExpandAggregatePeephole(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx);
}
|
