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
145
146
147
|
#pragma once
#include "proto_variant.h"
#include <yql/essentials/public/purecalc/io_specs/protobuf_raw/spec.h>
namespace NYql {
namespace NPureCalc {
/**
* Processing mode for working with non-raw protobuf messages.
*
* @tparam T message type.
*/
template <typename T>
class TProtobufInputSpec: public TProtobufRawInputSpec {
static_assert(std::is_base_of<google::protobuf::Message, T>::value,
"should be derived from google::protobuf::Message");
public:
TProtobufInputSpec(
const TMaybe<TString>& timestampColumn = Nothing(),
const TProtoSchemaOptions& options = {}
)
: TProtobufRawInputSpec(*T::descriptor(), timestampColumn, options)
{
}
};
/**
* Processing mode for working with non-raw protobuf messages.
*
* @tparam T message type.
*/
template <typename T>
class TProtobufOutputSpec: public TProtobufRawOutputSpec {
static_assert(std::is_base_of<google::protobuf::Message, T>::value,
"should be derived from google::protobuf::Message");
public:
TProtobufOutputSpec(
const TProtoSchemaOptions& options = {},
google::protobuf::Arena* arena = nullptr
)
: TProtobufRawOutputSpec(*T::descriptor(), nullptr, options, arena)
{
}
};
/**
* Processing mode for working with non-raw protobuf messages and several outputs.
*/
template <typename... T>
class TProtobufMultiOutputSpec: public TProtobufRawMultiOutputSpec {
static_assert(
std::conjunction_v<std::is_base_of<google::protobuf::Message, T>...>,
"all types should be derived from google::protobuf::Message");
public:
TProtobufMultiOutputSpec(
const TProtoSchemaOptions& options = {},
TMaybe<TVector<google::protobuf::Arena*>> arenas = {}
)
: TProtobufRawMultiOutputSpec({T::descriptor()...}, Nothing(), options, std::move(arenas))
{
}
};
template <typename T>
struct TInputSpecTraits<TProtobufInputSpec<T>> {
static const constexpr bool IsPartial = false;
static const constexpr bool SupportPullStreamMode = true;
static const constexpr bool SupportPullListMode = true;
static const constexpr bool SupportPushStreamMode = true;
using TConsumerType = THolder<IConsumer<T*>>;
static void PreparePullStreamWorker(const TProtobufInputSpec<T>& inputSpec, IPullStreamWorker* worker, THolder<IStream<T*>> stream) {
auto raw = ConvertStream<google::protobuf::Message*>(std::move(stream));
TInputSpecTraits<TProtobufRawInputSpec>::PreparePullStreamWorker(inputSpec, worker, std::move(raw));
}
static void PreparePullListWorker(const TProtobufInputSpec<T>& inputSpec, IPullListWorker* worker, THolder<IStream<T*>> stream) {
auto raw = ConvertStream<google::protobuf::Message*>(std::move(stream));
TInputSpecTraits<TProtobufRawInputSpec>::PreparePullListWorker(inputSpec, worker, std::move(raw));
}
static TConsumerType MakeConsumer(const TProtobufInputSpec<T>& inputSpec, TWorkerHolder<IPushStreamWorker> worker) {
auto raw = TInputSpecTraits<TProtobufRawInputSpec>::MakeConsumer(inputSpec, std::move(worker));
return ConvertConsumer<T*>(std::move(raw));
}
};
template <typename T>
struct TOutputSpecTraits<TProtobufOutputSpec<T>> {
static const constexpr bool IsPartial = false;
static const constexpr bool SupportPullStreamMode = true;
static const constexpr bool SupportPullListMode = true;
static const constexpr bool SupportPushStreamMode = true;
using TOutputItemType = T*;
using TPullStreamReturnType = THolder<IStream<TOutputItemType>>;
using TPullListReturnType = THolder<IStream<TOutputItemType>>;
static TPullStreamReturnType ConvertPullStreamWorkerToOutputType(const TProtobufOutputSpec<T>& outputSpec, TWorkerHolder<IPullStreamWorker> worker) {
auto raw = TOutputSpecTraits<TProtobufRawOutputSpec>::ConvertPullStreamWorkerToOutputType(outputSpec, std::move(worker));
return ConvertStreamUnsafe<TOutputItemType>(std::move(raw));
}
static TPullListReturnType ConvertPullListWorkerToOutputType(const TProtobufOutputSpec<T>& outputSpec, TWorkerHolder<IPullListWorker> worker) {
auto raw = TOutputSpecTraits<TProtobufRawOutputSpec>::ConvertPullListWorkerToOutputType(outputSpec, std::move(worker));
return ConvertStreamUnsafe<TOutputItemType>(std::move(raw));
}
static void SetConsumerToWorker(const TProtobufOutputSpec<T>& outputSpec, IPushStreamWorker* worker, THolder<IConsumer<T*>> consumer) {
auto raw = ConvertConsumerUnsafe<google::protobuf::Message*>(std::move(consumer));
TOutputSpecTraits<TProtobufRawOutputSpec>::SetConsumerToWorker(outputSpec, worker, std::move(raw));
}
};
template <typename... T>
struct TOutputSpecTraits<TProtobufMultiOutputSpec<T...>> {
static const constexpr bool IsPartial = false;
static const constexpr bool SupportPullStreamMode = true;
static const constexpr bool SupportPullListMode = true;
static const constexpr bool SupportPushStreamMode = true;
using TOutputItemType = std::variant<T*...>;
using TPullStreamReturnType = THolder<IStream<TOutputItemType>>;
using TPullListReturnType = THolder<IStream<TOutputItemType>>;
static TPullStreamReturnType ConvertPullStreamWorkerToOutputType(const TProtobufMultiOutputSpec<T...>& outputSpec, TWorkerHolder<IPullStreamWorker> worker) {
auto raw = TOutputSpecTraits<TProtobufRawMultiOutputSpec>::ConvertPullStreamWorkerToOutputType(outputSpec, std::move(worker));
return THolder(new NPrivate::TProtobufsMappingStream<T...>(std::move(raw)));
}
static TPullListReturnType ConvertPullListWorkerToOutputType(const TProtobufMultiOutputSpec<T...>& outputSpec, TWorkerHolder<IPullListWorker> worker) {
auto raw = TOutputSpecTraits<TProtobufRawMultiOutputSpec>::ConvertPullListWorkerToOutputType(outputSpec, std::move(worker));
return THolder(new NPrivate::TProtobufsMappingStream<T...>(std::move(raw)));
}
static void SetConsumerToWorker(const TProtobufMultiOutputSpec<T...>& outputSpec, IPushStreamWorker* worker, THolder<IConsumer<TOutputItemType>> consumer) {
auto wrapper = MakeHolder<NPrivate::TProtobufsMappingConsumer<T...>>(std::move(consumer));
TOutputSpecTraits<TProtobufRawMultiOutputSpec>::SetConsumerToWorker(outputSpec, worker, std::move(wrapper));
}
};
}
}
|
