Skip to content

Commit 16918da

Browse files
igormunkinigormunkin
committed
Fix TArrowBlock::From and its usage
commit_hash:0d0fe61d962e4c365f99bd84ebdb2229696c4dc0
1 parent 23c0506 commit 16918da

File tree

8 files changed

+53
-28
lines changed

8 files changed

+53
-28
lines changed

‎yql/essentials/minikql/comp_nodes/mkql_apply.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ class TApplyWrapper: public TMutableCodegeneratorPtrNode<TApplyWrapper> {
5454
state.Args[i] = state.HolderFactory.CreateArrowBlock(arrow::Datum(batch.values[i]));
5555
}
5656

57-
const auto ret = Callable_.Run(&state.ValueBuilder, state.Args.data());
57+
const auto& ret = Callable_.Run(&state.ValueBuilder, state.Args.data());
5858
*res = TArrowBlock::From(ret).GetDatum();
5959
return arrow::Status::OK();
6060
})

‎yql/essentials/minikql/comp_nodes/mkql_block_compress.cpp

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,9 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TCompressWithScalarB
8484

8585
const auto bitmapValue = getres.second[BitmapIndex_](ctx, block);
8686
const auto bitmap = CallInst::Create(getBitmap, { WrapArgumentForWindows(bitmapValue, ctx, block) }, "bitmap", block);
87+
88+
ValueCleanup(EValueRepresentation::Any, bitmapValue, ctx, block);
89+
8790
const auto one = ConstantInt::get(bitmapType, 1);
8891
const auto band = BinaryOperator::CreateAnd(bitmap, one, "band", block);
8992
const auto good = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, band, one, "good", block);
@@ -189,6 +192,9 @@ using TBaseComputation = TStatelessWideFlowCodegeneratorNode<TCompressScalars>;
189192

190193
const auto bitmapValue = getres.second[BitmapIndex_](ctx, block);
191194
const auto pops = CallInst::Create(getPopCount, { WrapArgumentForWindows(bitmapValue, ctx, block) }, "pops", block);
195+
196+
ValueCleanup(EValueRepresentation::Any, bitmapValue, ctx, block);
197+
192198
const auto good = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_UGT, pops, ConstantInt::get(sizeType, 0), "good", block);
193199

194200
BranchInst::Create(fill, loop, good, block);
@@ -270,7 +276,7 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TCompressBlocks>;
270276
s.IsFinished_ = true;
271277
break;
272278
case EFetchResult::One:
273-
switch (s.Check(bitmap.Release())) {
279+
switch (s.Check(bitmap)) {
274280
case TState::EStep::Copy:
275281
for (ui32 i = 0; i < s.Values.size(); ++i) {
276282
if (const auto out = output[i]) {
@@ -415,6 +421,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TCompressBlocks>;
415421
const auto checkPtr = CastInst::Create(Instruction::IntToPtr, checkFunc, PointerType::getUnqual(checkType), "check_func", block);
416422
const auto check = CallInst::Create(checkType, checkPtr, {stateArg, bitmapArg}, "check", block);
417423

424+
ValueCleanup(EValueRepresentation::Any, bitmap, ctx, block);
425+
418426
result->addIncoming(ConstantInt::get(statusType, static_cast<i32>(EFetchResult::One)), block);
419427

420428
const auto step = SwitchInst::Create(check, save, 2U, block);
@@ -552,9 +560,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TCompressBlocks>;
552560
EStep Check(const NUdf::TUnboxedValuePod bitmapValue) {
553561
Y_ABORT_UNLESS(!IsFinished_);
554562
Y_ABORT_UNLESS(!InputSize_);
555-
const NUdf::TUnboxedValue b(std::move(bitmapValue));
556563
auto& bitmap = Arrays_.back();
557-
bitmap = TArrowBlock::From(b).GetDatum().array();
564+
bitmap = TArrowBlock::From(bitmapValue).GetDatum().array();
558565

559566
if (!bitmap->length)
560567
return EStep::Skip;

‎yql/essentials/minikql/comp_nodes/mkql_block_skiptake.cpp

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,12 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideSkipBlocksWrapp
117117

118118

119119
block = test;
120-
const auto height = CallInst::Create(getCount, { WrapArgumentForWindows(getres.second.back()(ctx, block), ctx, block) }, "height", block);
120+
121+
const auto countValue = getres.second.back()(ctx, block);
122+
const auto height = CallInst::Create(getCount, { WrapArgumentForWindows(countValue, ctx, block) }, "height", block);
123+
124+
ValueCleanup(EValueRepresentation::Any, countValue, ctx, block);
125+
121126
const auto part = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_ULT, count, height, "part", block);
122127
const auto decr = BinaryOperator::CreateSub(count, height, "decr", block);
123128
count->addIncoming(decr, block);
@@ -201,6 +206,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideSkipBlocksWrapp
201206

202207
const auto slice = CallInst::Create(sliceType, slicePtr, {ctx.GetFactory(), value, offset}, "slice", block);
203208

209+
ValueCleanup(EValueRepresentation::Any, value, ctx, block);
210+
204211
output->addIncoming(slice, block);
205212
BranchInst::Create(exit, block);
206213

@@ -214,9 +221,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideSkipBlocksWrapp
214221
#endif
215222
private:
216223
static NUdf::TUnboxedValuePod SliceBlock(const THolderFactory& holderFactory, NUdf::TUnboxedValuePod block, const uint64_t offset) {
217-
NUdf::TUnboxedValue b(block);
218-
auto& datum = TArrowBlock::From(b).GetDatum();
219-
return datum.is_scalar() ? b.Release() : holderFactory.CreateArrowBlock(DeepSlice(datum.array(), offset, datum.array()->length - offset));
224+
const auto& datum = TArrowBlock::From(block).GetDatum();
225+
return datum.is_scalar() ? block : holderFactory.CreateArrowBlock(DeepSlice(datum.array(), offset, datum.array()->length - offset));
220226
}
221227

222228
void RegisterDependencies() const final {
@@ -325,7 +331,11 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeBlocksWrapp
325331

326332
block = good;
327333

328-
const auto height = CallInst::Create(getCount, { WrapArgumentForWindows(getres.second.back()(ctx, block), ctx, block) }, "height", block);
334+
const auto countValue = getres.second.back()(ctx, block);
335+
const auto height = CallInst::Create(getCount, { WrapArgumentForWindows(countValue, ctx, block) }, "height", block);
336+
337+
ValueCleanup(EValueRepresentation::Any, countValue, ctx, block);
338+
329339
const auto part = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_ULT, count, height, "part", block);
330340
const auto decr = BinaryOperator::CreateSub(count, height, "decr", block);
331341

@@ -395,6 +405,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeBlocksWrapp
395405

396406
const auto slice = CallInst::Create(sliceType, slicePtr, {ctx.GetFactory(), value, size}, "slice", block);
397407

408+
ValueCleanup(EValueRepresentation::Any, value, ctx, block);
409+
398410
output->addIncoming(slice, block);
399411
BranchInst::Create(exit, block);
400412

@@ -408,9 +420,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeBlocksWrapp
408420
#endif
409421
private:
410422
static NUdf::TUnboxedValuePod SliceBlock(const THolderFactory& holderFactory, NUdf::TUnboxedValuePod block, const uint64_t offset) {
411-
NUdf::TUnboxedValue b(block);
412-
auto& datum = TArrowBlock::From(b).GetDatum();
413-
return datum.is_scalar() ? b.Release() : holderFactory.CreateArrowBlock(DeepSlice(datum.array(), 0ULL, offset));
423+
const auto& datum = TArrowBlock::From(block).GetDatum();
424+
return datum.is_scalar() ? block : holderFactory.CreateArrowBlock(DeepSlice(datum.array(), 0ULL, offset));
414425
}
415426

416427
void RegisterDependencies() const final {

‎yql/essentials/minikql/comp_nodes/mkql_blocks.cpp

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -383,7 +383,7 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TFromBlocksWrapper>;
383383
if (auto item = s.GetValue(ctx.HolderFactory); !item.IsInvalid())
384384
return item;
385385

386-
if (const auto input = Flow_->GetValue(ctx).Release(); input.IsSpecial())
386+
if (const auto input = Flow_->GetValue(ctx); input.IsSpecial())
387387
return input;
388388
else
389389
s.Reset(input);
@@ -443,6 +443,8 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TFromBlocksWrapper>;
443443
const auto setPtr = CastInst::Create(Instruction::IntToPtr, setFunc, PointerType::getUnqual(setType), "set", block);
444444
CallInst::Create(setType, setPtr, {stateArg, input }, "", block);
445445

446+
ValueCleanup(EValueRepresentation::Any, input, ctx, block);
447+
446448
BranchInst::Create(work, block);
447449

448450
block = done;
@@ -475,8 +477,7 @@ using TBaseComputation = TStatefulFlowCodegeneratorNode<TFromBlocksWrapper>;
475477
}
476478

477479
void Reset(const NUdf::TUnboxedValuePod block) {
478-
const NUdf::TUnboxedValue v(block);
479-
const auto& datum = TArrowBlock::From(v).GetDatum();
480+
const auto& datum = TArrowBlock::From(block).GetDatum();
480481
MKQL_ENSURE(datum.is_arraylike(), "Expecting array as FromBlocks argument");
481482
MKQL_ENSURE(Arrays_.empty(), "Not all input is processed");
482483
if (datum.is_array()) {
@@ -634,6 +635,8 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideFromBlocksWrapp
634635
const auto countValue = getres.second.back()(ctx, block);
635636
const auto height = CallInst::Create(getCount, { WrapArgumentForWindows(countValue, ctx, block) }, "height", block);
636637

638+
ValueCleanup(EValueRepresentation::Any, countValue, ctx, block);
639+
637640
new StoreInst(height, countPtr, block);
638641
new StoreInst(ConstantInt::get(indexType, 0), indexPtr, block);
639642

@@ -953,9 +956,8 @@ using TBaseComputation = TMutableCodegeneratorNode<TReplicateScalarWrapper>;
953956
}
954957

955958
arrow::Datum DoReplicate(const NUdf::TUnboxedValuePod val, const NUdf::TUnboxedValuePod cnt, TComputationContext& ctx) const {
956-
const NUdf::TUnboxedValue v(val), c(cnt);
957-
const auto value = TArrowBlock::From(v).GetDatum().scalar();
958-
const ui64 count = TArrowBlock::From(c).GetDatum().scalar_as<arrow::UInt64Scalar>().value;
959+
const auto value = TArrowBlock::From(val).GetDatum().scalar();
960+
const ui64 count = TArrowBlock::From(cnt).GetDatum().scalar_as<arrow::UInt64Scalar>().value;
959961

960962
const auto reader = MakeBlockReader(TTypeInfoHelper(), Type_);
961963
const auto builder = MakeArrayBuilder(TTypeInfoHelper(), Type_, ctx.ArrowMemoryPool, count, &ctx.Builder->GetPgBuilder());

‎yql/essentials/minikql/computation/mkql_block_impl.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@ extern "C" uint64_t GetBlockCount(const NYql::NUdf::TUnboxedValuePod data) {
1818
}
1919

2020
extern "C" uint64_t GetBitmapPopCountCount(const NYql::NUdf::TUnboxedValuePod data) {
21-
const NYql::NUdf::TUnboxedValue v(data);
22-
const auto& arr = NKikimr::NMiniKQL::TArrowBlock::From(v).GetDatum().array();
21+
const auto& arr = NKikimr::NMiniKQL::TArrowBlock::From(data).GetDatum().array();
2322
const size_t len = (size_t)arr->length;
2423
MKQL_ENSURE(arr->GetNullCount() == 0, "Bitmap block should not have nulls");
2524
const ui8* src = arr->GetValues<ui8>(1);
@@ -259,7 +258,8 @@ NUdf::TUnboxedValuePod TBlockFuncNode::DoCalculate(TComputationContext& ctx) con
259258

260259
std::vector<arrow::Datum> argDatums;
261260
for (ui32 i = 0; i < ArgsNodes.size(); ++i) {
262-
argDatums.emplace_back(TArrowBlock::From(ArgsNodes[i]->GetValue(ctx)).GetDatum());
261+
const auto& value = ArgsNodes[i]->GetValue(ctx);
262+
argDatums.emplace_back(TArrowBlock::From(value).GetDatum());
263263
ARROW_DEBUG_CHECK_DATUM_TYPES(ArgsValuesDescr[i], argDatums.back().descr());
264264
}
265265

@@ -352,7 +352,7 @@ void TBlockState::FillArrays() {
352352

353353
for (size_t i = 0U; i < Deques.size(); ++i) {
354354
Deques[i].clear();
355-
if (const auto value = Values[i]) {
355+
if (const auto& value = Values[i]) {
356356
const auto& datum = TArrowBlock::From(value).GetDatum();
357357
if (datum.is_scalar()) {
358358
return;

‎yql/essentials/minikql/computation/mkql_computation_node.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@ TDatumProvider MakeDatumProvider(const arrow::Datum& datum) {
3939

4040
TDatumProvider MakeDatumProvider(const IComputationNode* node, TComputationContext& ctx) {
4141
return [node, &ctx]() {
42-
return TArrowBlock::From(node->GetValue(ctx)).GetDatum();
42+
const auto& value = node->GetValue(ctx);
43+
return TArrowBlock::From(value).GetDatum();
4344
};
4445
}
4546

‎yql/essentials/minikql/computation/mkql_computation_node_holders.h

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -582,11 +582,13 @@ class TArrowBlock: public TComputationValue<TArrowBlock> {
582582
{
583583
}
584584

585-
inline static TArrowBlock& From(const NUdf::TUnboxedValue& value) {
586-
return *static_cast<TArrowBlock*>(value.AsBoxed().Get());
585+
inline static const TArrowBlock& From(const NUdf::TUnboxedValuePod& value) {
586+
return *static_cast<TArrowBlock*>(value.AsRawBoxed());
587587
}
588588

589-
inline arrow::Datum& GetDatum() {
589+
inline static const TArrowBlock& From(NUdf::TUnboxedValuePod&& value) = delete;
590+
591+
inline const arrow::Datum& GetDatum() const {
590592
return Datum_;
591593
}
592594

‎yql/essentials/public/purecalc/io_specs/arrow/spec.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,8 @@ class TArrowOutputConverter {
170170
OutputItemType batch = Batch_.Get();
171171
size_t nvalues = DatumToMemberIDMap_.size();
172172

173-
const auto& sizeDatum = TArrowBlock::From(value.GetElement(BatchLengthID_)).GetDatum();
173+
const auto& sizeValue = value.GetElement(BatchLengthID_);
174+
const auto& sizeDatum = TArrowBlock::From(sizeValue).GetDatum();
174175
Y_ENSURE(sizeDatum.is_scalar());
175176
const auto& sizeScalar = sizeDatum.scalar();
176177
const auto& sizeData = arrow::internal::checked_cast<const arrow::UInt64Scalar&>(*sizeScalar);
@@ -179,7 +180,8 @@ class TArrowOutputConverter {
179180
TVector<arrow::Datum> datums(nvalues);
180181
for (size_t i = 0; i < nvalues; i++) {
181182
const ui32 id = DatumToMemberIDMap_[i];
182-
const auto& datum = TArrowBlock::From(value.GetElement(id)).GetDatum();
183+
const auto& datumValue = value.GetElement(id);
184+
const auto& datum = TArrowBlock::From(datumValue).GetDatum();
183185
datums[i] = datum;
184186
if (datum.is_scalar()) {
185187
continue;

0 commit comments

Comments
 (0)