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
|
#include "mkql_block_guess.h"
#include <yql/essentials/minikql/computation/mkql_block_impl.h>
#include <yql/essentials/minikql/computation/mkql_block_reader.h>
#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
#include <yql/essentials/minikql/mkql_node_builder.h>
#include <yql/essentials/minikql/mkql_node_cast.h>
#include <yql/essentials/public/udf/arrow/block_builder.h>
#include <yql/essentials/public/udf/arrow/block_reader.h>
namespace NKikimr {
namespace NMiniKQL {
namespace {
template <bool IsOptional>
class TGuessBlockExec {
public:
class TGuessKernelState: public arrow::compute::KernelState {
public:
explicit TGuessKernelState(TType* inputItemType)
: Reader_(MakeBlockReader(TTypeInfoHelper(), inputItemType))
{
}
IBlockReader& GetReader() {
return *Reader_;
}
private:
std::unique_ptr<IBlockReader> Reader_;
};
explicit TGuessBlockExec(TType* inputItemType, ui32 alternativeIndex, TType* resultItemType)
: InputItemType_(inputItemType)
, AlternativeIndex_(alternativeIndex)
, ResultItemType_(resultItemType)
{
}
arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) const {
auto& reader = static_cast<TGuessKernelState&>(*ctx->state()).GetReader();
const arrow::Datum& variantDatum = batch.values[0];
if (variantDatum.is_scalar()) {
*res = ConvertScalar(ResultItemType_,
ComputeOutputItem(reader.GetScalarItem(*variantDatum.scalar())),
*ctx->memory_pool());
return arrow::Status::OK();
}
MKQL_ENSURE(variantDatum.is_array(), "Expected array datum");
const auto& variantArrayData = variantDatum.array();
const size_t length = static_cast<size_t>(variantArrayData->length);
auto builder = NYql::NUdf::MakeArrayBuilder(TTypeInfoHelper(), ResultItemType_, *ctx->memory_pool(), length, /*pgBuilder=*/nullptr);
for (size_t i = 0; i < length; ++i) {
builder->Add(ComputeOutputItem(reader.GetItem(*variantArrayData, i)));
}
*res = builder->Build(/*finish=*/true);
return arrow::Status::OK();
}
private:
TBlockItem ComputeOutputItem(TBlockItem blockItem) const {
if constexpr (IsOptional) {
if (!blockItem) {
return TBlockItem{};
}
}
return blockItem.GetVariantIndex() == AlternativeIndex_
? blockItem.GetVariantItem().MakeOptional()
: TBlockItem{};
}
TType* const InputItemType_;
const ui32 AlternativeIndex_;
TType* const ResultItemType_;
};
template <bool IsOptional>
std::shared_ptr<arrow::compute::ScalarKernel> MakeBlockGuessKernel(const TVector<TType*>& argTypes,
TType* resultType,
TType* inputItemType,
ui32 alternativeIndex) {
using TExec = TGuessBlockExec<IsOptional>;
auto exec = std::make_shared<TExec>(
inputItemType,
alternativeIndex,
AS_TYPE(TBlockType, resultType)->GetItemType());
auto kernel = std::make_shared<arrow::compute::ScalarKernel>(
ConvertToInputTypes(argTypes),
ConvertToOutputType(resultType),
[exec](arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
return exec->Exec(ctx, batch, res);
});
kernel->null_handling = arrow::compute::NullHandling::COMPUTED_NO_PREALLOCATE;
kernel->mem_allocation = arrow::compute::MemAllocation::NO_PREALLOCATE;
kernel->init = [inputItemType](arrow::compute::KernelContext*, const arrow::compute::KernelInitArgs&) {
return arrow::Result(std::make_unique<typename TExec::TGuessKernelState>(inputItemType));
};
return kernel;
}
} // namespace
IComputationNode* WrapBlockGuess(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
MKQL_ENSURE(callable.GetInputsCount() == 2, "Expected 2 arguments");
auto blockType = AS_TYPE(TBlockType, callable.GetInput(0).GetStaticType());
auto inputItemType = blockType->GetItemType();
bool isOptional;
auto variantItemType = UnpackOptional(inputItemType, isOptional);
auto variantType = AS_TYPE(TVariantType, variantItemType);
const ui32 alternativeIndex = AS_VALUE(TDataLiteral, callable.GetInput(1))->AsValue().Get<ui32>();
MKQL_ENSURE(alternativeIndex < variantType->GetAlternativesCount(), "Bad alternative index");
auto variantCompute = LocateNode(ctx.NodeLocator, callable, 0);
TComputationNodePtrVector argsNodes = {variantCompute};
TVector<TType*> argsTypes = {blockType};
auto resultType = callable.GetType()->GetReturnType();
auto kernel = isOptional
? MakeBlockGuessKernel<true>(argsTypes, resultType, inputItemType, alternativeIndex)
: MakeBlockGuessKernel<false>(argsTypes, resultType, inputItemType, alternativeIndex);
return new TBlockFuncNode(ctx.Mutables, ctx.RuntimeSettings->DatumValidation.Get(),
callable.GetType()->GetName(), std::move(argsNodes), argsTypes, resultType, *kernel, kernel);
}
} // namespace NMiniKQL
} // namespace NKikimr
|
