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
|
#include "mkql_block_variant.h"
#include <yql/essentials/minikql/computation/mkql_block_impl.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/minikql/mkql_type_builder.h>
#include <yql/essentials/public/udf/arrow/dense_union_scalar.h>
#include <yql/essentials/public/udf/arrow/util.h>
#include <numeric>
namespace NKikimr::NMiniKQL {
namespace {
class TVariantBlockExec {
public:
TVariantBlockExec(ui32 alternativeIndex, std::shared_ptr<arrow::DataType> resultArrowType,
TVector<std::shared_ptr<arrow::DataType>> alternativeArrowTypes)
: AlternativeIndex_(alternativeIndex)
, ResultArrowType_(std::move(resultArrowType))
, AlternativeArrowTypes_(std::move(alternativeArrowTypes))
{
}
arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) const {
const arrow::Datum& payloadDatum = batch.values[0];
if (payloadDatum.is_scalar()) {
*res = arrow::Datum(std::make_shared<NYql::NUdf::TDenseUnionScalar>(payloadDatum.scalar(), AlternativeIndex_, ResultArrowType_));
return arrow::Status::OK();
}
MKQL_ENSURE(payloadDatum.is_array(), "Expected array datum");
*res = AttachAsVariantChild(payloadDatum.array(), ctx->memory_pool());
return arrow::Status::OK();
}
private:
std::shared_ptr<arrow::ArrayData> AttachAsVariantChild(const std::shared_ptr<arrow::ArrayData>& payload, arrow::MemoryPool* pool) const {
const i64 length = payload->length;
auto typeCodes = ARROW_RESULT(arrow::AllocateBuffer(length, pool));
std::fill_n(reinterpret_cast<i8*>(typeCodes->mutable_data()), length, static_cast<i8>(AlternativeIndex_));
auto valueOffsets = ARROW_RESULT(arrow::AllocateBuffer(length * static_cast<i64>(sizeof(i32)), pool));
auto* offsets = reinterpret_cast<i32*>(valueOffsets->mutable_data());
std::iota(offsets, offsets + length, 0);
std::vector<std::shared_ptr<arrow::ArrayData>> children(AlternativeArrowTypes_.size());
for (size_t i = 0; i < children.size(); ++i) {
children[i] = (i == AlternativeIndex_) ? payload : NYql::NUdf::MakeEmptyArray(AlternativeArrowTypes_[i], pool);
}
return arrow::ArrayData::Make(ResultArrowType_, length,
{nullptr, std::move(typeCodes), std::move(valueOffsets)},
std::move(children), /*null_count=*/0, /*offset=*/0);
}
const ui32 AlternativeIndex_;
const std::shared_ptr<arrow::DataType> ResultArrowType_;
const TVector<std::shared_ptr<arrow::DataType>> AlternativeArrowTypes_;
};
TVector<std::shared_ptr<arrow::DataType>> ConvertAlternativeArrowTypes(const TVariantType& variantType) {
TVector<std::shared_ptr<arrow::DataType>> alternativeArrowTypes(variantType.GetAlternativesCount());
for (ui32 i = 0; i < alternativeArrowTypes.size(); ++i) {
MKQL_ENSURE(ConvertArrowType(variantType.GetAlternativeType(i), alternativeArrowTypes[i]), "Unsupported arrow type");
}
return alternativeArrowTypes;
}
std::shared_ptr<arrow::compute::ScalarKernel> MakeBlockVariantKernel(const TVector<TType*>& argTypes,
TType* resultType,
ui32 alternativeIndex,
const TVariantType& variantType) {
TType* resultItemType = AS_TYPE(TBlockType, resultType)->GetItemType();
std::shared_ptr<arrow::DataType> resultArrowType;
MKQL_ENSURE(ConvertArrowType(resultItemType, resultArrowType), "Unsupported arrow type");
auto exec = std::make_shared<TVariantBlockExec>(
alternativeIndex, std::move(resultArrowType), ConvertAlternativeArrowTypes(variantType));
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;
return kernel;
}
} // namespace
IComputationNode* WrapBlockVariant(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
MKQL_ENSURE(callable.GetInputsCount() == 2, "Expected 2 arguments");
auto payloadBlockType = AS_TYPE(TBlockType, callable.GetInput(0).GetStaticType());
const ui32 alternativeIndex = AS_VALUE(TDataLiteral, callable.GetInput(1))->AsValue().Get<ui32>();
auto resultType = callable.GetType()->GetReturnType();
auto variantItemType = AS_TYPE(TBlockType, resultType)->GetItemType();
auto variantType = AS_TYPE(TVariantType, variantItemType);
MKQL_ENSURE(alternativeIndex < variantType->GetAlternativesCount(), "Bad alternative index");
auto payloadCompute = LocateNode(ctx.NodeLocator, callable, 0);
TComputationNodePtrVector argsNodes = {payloadCompute};
TVector<TType*> argsTypes = {payloadBlockType};
auto kernel = MakeBlockVariantKernel(argsTypes, resultType, alternativeIndex, *variantType);
return new TBlockFuncNode(ctx.Mutables, ctx.RuntimeSettings->DatumValidation.Get(),
callable.GetType()->GetName(), std::move(argsNodes), argsTypes, resultType, *kernel, kernel);
}
} // namespace NKikimr::NMiniKQL
|
