Skip to content

Commit 72cc40a

Browse files
avevadavevad
authored
Merge 1a26206 into bc843db 
2 parents bc843db + 1a26206 commit 72cc40a

File tree

5 files changed

+77
-58
lines changed

5 files changed

+77
-58
lines changed

‎ydb/library/yql/minikql/comp_nodes/mkql_extend.cpp‎

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -61,29 +61,37 @@ class TState : public TComputationValue<TState> {
6161
};
6262
#endif
6363

64-
class TExtendWideFlowWrapper : public TStatefulWideFlowCodegeneratorNode<TExtendWideFlowWrapper> {
64+
class TExtendWideFlowWrapper : public TStatefulWideFlowCodegeneratorNodeImpl<TExtendWideFlowWrapper, TState, true> {
6565
using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TExtendWideFlowWrapper>;
6666
public:
6767
TExtendWideFlowWrapper(TComputationMutables& mutables, TComputationWideFlowNodePtrVector&& flows, size_t width)
68-
: TBaseComputation(mutables, this, EValueRepresentation::Boxed)
68+
: TStatefulWideFlowCodegeneratorNodeImpl(mutables, this)
6969
, Flows_(std::move(flows)), Width_(width)
7070
{
7171
#ifdef MKQL_DISABLE_CODEGEN
7272
Y_UNUSED(Width_);
7373
#endif
7474
}
7575

76-
EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
77-
auto& s = GetState(state, ctx);
78-
while (s.Index >= 0) {
79-
switch (Flows_[s.Index]->FetchValues(ctx, output)) {
76+
void InitState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
77+
state = ctx.HolderFactory.Create<TState>(Flows_.size());
78+
}
79+
80+
// Codegen compatibility
81+
void MakeState(TComputationContext& ctx, NUdf::TUnboxedValue& state) const {
82+
InitState(state, ctx);
83+
}
84+
85+
TMaybe<EFetchResult> DoFetch(TState &state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
86+
while (state.Index >= 0) {
87+
switch (Flows_[state.Index]->FetchValues(ctx, output)) {
8088
case EFetchResult::One:
8189
return EFetchResult::One;
8290
case EFetchResult::Yield:
83-
s.NextFlow();
91+
state.NextFlow();
8492
return EFetchResult::Yield;
8593
case EFetchResult::Finish:
86-
s.FlowOver();
94+
state.FlowOver();
8795
break;
8896
}
8997
}
@@ -204,16 +212,6 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TExtendWideFlowWrapp
204212
}
205213
}
206214

207-
void MakeState(TComputationContext& ctx, NUdf::TUnboxedValue& state) const {
208-
state = ctx.HolderFactory.Create<TState>(Flows_.size());
209-
}
210-
211-
TState& GetState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
212-
if (!state.HasValue())
213-
MakeState(ctx, state);
214-
return *static_cast<TState*>(state.AsBoxed().Get());
215-
}
216-
217215
const TComputationWideFlowNodePtrVector Flows_;
218216
const size_t Width_;
219217
};

‎ydb/library/yql/minikql/comp_nodes/mkql_if.cpp‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -167,22 +167,21 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TFlowIfWrapper<IsOptiona
167167
IComputationNode* const ElseBranch;
168168
};
169169

170-
class TWideIfWrapper : public TStatefulWideFlowCodegeneratorNode<TWideIfWrapper> {
171-
using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideIfWrapper>;
170+
class TWideIfWrapper : public TStatefulWideFlowCodegeneratorNodeImpl<TWideIfWrapper, bool>{
172171
public:
173172
TWideIfWrapper(TComputationMutables& mutables, IComputationNode* predicate, IComputationWideFlowNode* thenBranch, IComputationWideFlowNode* elseBranch)
174-
: TBaseComputation(mutables, nullptr, EValueRepresentation::Embedded)
173+
: TStatefulWideFlowCodegeneratorNodeImpl<TWideIfWrapper, bool>(mutables, nullptr)
175174
, Predicate(predicate)
176175
, ThenBranch(thenBranch)
177176
, ElseBranch(elseBranch)
178177
{}
179178

180-
EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
181-
if (state.IsInvalid()) {
182-
state = Predicate->GetValue(ctx);
183-
}
179+
void InitState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
180+
state = Predicate->GetValue(ctx);
181+
}
184182

185-
return (state.Get<bool>() ? ThenBranch : ElseBranch)->FetchValues(ctx, output);
183+
TMaybe<EFetchResult> DoFetch(bool pred, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
184+
return (pred ? ThenBranch : ElseBranch)->FetchValues(ctx, output);
186185
}
187186
#ifndef MKQL_DISABLE_CODEGEN
188187
TGenerateResult DoGenGetValues(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {

‎ydb/library/yql/minikql/comp_nodes/mkql_skip.cpp‎

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -117,31 +117,25 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TSkipFlowWrapper>;
117117
IComputationNode* const Count;
118118
};
119119

120-
class TWideSkipWrapper : public TStatefulWideFlowCodegeneratorNode<TWideSkipWrapper> {
121-
using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideSkipWrapper>;
120+
class TWideSkipWrapper : public TStatefulWideFlowCodegeneratorNodeImpl<TWideSkipWrapper, ui64> {
122121
public:
123122
TWideSkipWrapper(TComputationMutables& mutables, IComputationWideFlowNode* flow, IComputationNode* count, ui32 size)
124-
: TBaseComputation(mutables, flow, EValueRepresentation::Embedded)
123+
: TStatefulWideFlowCodegeneratorNodeImpl<TWideSkipWrapper, ui64>(mutables, flow)
125124
, Flow(flow)
126125
, Count(count)
127126
, StubsIndex(mutables.IncrementWideFieldsIndex(size))
128127
{}
129128

130-
EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
131-
if (state.IsInvalid()) {
132-
state = Count->GetValue(ctx);
133-
}
134-
135-
if (auto count = state.Get<ui64>()) {
136-
do if (const auto result = Flow->FetchValues(ctx, ctx.WideFields.data() + StubsIndex); EFetchResult::One != result) {
137-
state = NUdf::TUnboxedValuePod(count);
138-
return result;
139-
} while (--count);
129+
void InitState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
130+
state = Count->GetValue(ctx);
131+
}
140132

141-
state = NUdf::TUnboxedValuePod::Zero();
133+
TMaybe<EFetchResult> DoFetch(ui64 &count, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
134+
if (auto res = Flow->FetchValues(ctx, output); res != EFetchResult::One || count == 0) {
135+
return res;
142136
}
143-
144-
return Flow->FetchValues(ctx, output);
137+
count--;
138+
return Nothing();
145139
}
146140

147141
#ifndef MKQL_DISABLE_CODEGEN

‎ydb/library/yql/minikql/comp_nodes/mkql_take.cpp‎

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -95,28 +95,22 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TTakeFlowWrapper>;
9595
IComputationNode* const Count;
9696
};
9797

98-
class TWideTakeWrapper : public TStatefulWideFlowCodegeneratorNode<TWideTakeWrapper> {
99-
using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeWrapper>;
98+
class TWideTakeWrapper : public TStatefulWideFlowCodegeneratorNodeImpl<TWideTakeWrapper, ui64> {
10099
public:
101100
TWideTakeWrapper(TComputationMutables& mutables, IComputationWideFlowNode* flow, IComputationNode* count)
102-
: TBaseComputation(mutables, flow, EValueRepresentation::Embedded), Flow(flow), Count(count)
101+
: TStatefulWideFlowCodegeneratorNodeImpl<TWideTakeWrapper, ui64>(mutables, flow), Flow(flow), Count(count)
103102
{}
104103

105-
EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
106-
if (state.IsInvalid()) {
107-
state = Count->GetValue(ctx);
108-
}
104+
void InitState(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
105+
state = Count->GetValue(ctx);
106+
}
109107

110-
if (auto count = state.Get<ui64>()) {
111-
if (const auto result = Flow->FetchValues(ctx, output); EFetchResult::One == result) {
112-
state = NUdf::TUnboxedValuePod(--count);
113-
return EFetchResult::One;
114-
} else {
115-
return result;
116-
}
108+
TMaybe<EFetchResult> DoFetch(ui64 &count, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
109+
if (count == 0) {
110+
return EFetchResult::Finish;
117111
}
118-
119-
return EFetchResult::Finish;
112+
count--;
113+
return Flow->FetchValues(ctx, output);
120114
}
121115
#ifndef MKQL_DISABLE_CODEGEN
122116
TGenerateResult DoGenGetValues(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {

‎ydb/library/yql/minikql/computation/mkql_computation_node_codegen.h.txt‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1486,3 +1486,37 @@ Value* CheckAdjustedMemLimit(ui64 limit, Value* init, const TCodegenContext& ctx
14861486
}
14871487
}
14881488
#endif
1489+
1490+
namespace NKikimr::NMiniKQL {
1491+
1492+
template<typename TDerivedWrapper, typename TState, bool BoxedState = false>
1493+
class TStatefulWideFlowCodegeneratorNodeImpl : public TStatefulWideFlowCodegeneratorNode<TDerivedWrapper> {
1494+
public:
1495+
using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TDerivedWrapper>;
1496+
1497+
TStatefulWideFlowCodegeneratorNodeImpl(TComputationMutables &mutables, IComputationWideFlowNode *flow)
1498+
: TBaseComputation(mutables, flow, BoxedState ? EValueRepresentation::Boxed : EValueRepresentation::Embedded) {}
1499+
1500+
EFetchResult
1501+
DoCalculate(NUdf::TUnboxedValue &state, TComputationContext &ctx, NUdf::TUnboxedValue *const *output) const {
1502+
if (state.IsInvalid()) {
1503+
static_cast<const TDerivedWrapper *>(this)->InitState(state, ctx);
1504+
}
1505+
1506+
TState *ptr = nullptr;
1507+
if constexpr (BoxedState) {
1508+
ptr = static_cast<TState *>(state.AsBoxed().Get());
1509+
} else {
1510+
ptr = static_cast<TState *>(state.GetRawPtr());
1511+
}
1512+
1513+
while (true) {
1514+
TMaybe<EFetchResult> result = static_cast<const TDerivedWrapper *>(this)->DoFetch(*ptr, ctx, output);
1515+
if (!result.Empty()) {
1516+
return result.GetRef();
1517+
}
1518+
}
1519+
}
1520+
};
1521+
1522+
}

0 commit comments

Comments
 (0)