Skip to content

Commit 0996aee

Browse files
evanevanevanevannnnevanevanevanevannnn
authored
Merge c7a72cb into 010b521 
2 parents 010b521 + c7a72cb commit 0996aee

File tree

7 files changed

+130
-23
lines changed

7 files changed

+130
-23
lines changed

‎ydb/core/external_sources/object_storage.cpp‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -333,7 +333,7 @@ struct TObjectStorageExternalSource : public IExternalSource {
333333
}
334334
for (const auto& entry : entries.Objects) {
335335
if (entry.Size > 0) {
336-
return entry.Path;
336+
return entry;
337337
}
338338
}
339339
throw yexception() << "couldn't find any files for type inference, please check that the right path is provided";
@@ -350,9 +350,9 @@ struct TObjectStorageExternalSource : public IExternalSource {
350350

351351
auto fileFormat = NObjectStorage::NInference::ConvertFileFormat(*format);
352352
auto arrowFetcherId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowFetchingActor(s3FetcherId, fileFormat, meta->Attributes));
353-
auto arrowInferencinatorId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowInferencinator(arrowFetcherId, fileFormat, meta->Attributes));
353+
auto arrowInferencinatorId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowInferencinator(arrowFetcherId, s3FetcherId, fileFormat, meta->Attributes));
354354

355-
return afterListing.Apply([arrowInferencinatorId, meta, actorSystem = ActorSystem](const NThreading::TFuture<TString>& pathFut) {
355+
return afterListing.Apply([arrowInferencinatorId, meta, actorSystem = ActorSystem](const NThreading::TFuture<NYql::NS3Lister::TObjectListEntry>& entryFut) {
356356
auto promise = NThreading::NewPromise<TMetadataResult>();
357357
auto schemaToMetadata = [meta](NThreading::TPromise<TMetadataResult> metaPromise, NObjectStorage::TEvInferredFileSchema&& response) {
358358
if (!response.Status.IsSuccess()) {
@@ -370,9 +370,10 @@ struct TObjectStorageExternalSource : public IExternalSource {
370370
result.Metadata = meta;
371371
metaPromise.SetValue(std::move(result));
372372
};
373+
auto [path, size, _] = entryFut.GetValue();
373374
actorSystem->Register(new NKqp::TActorRequestHandler<NObjectStorage::TEvInferFileSchema, NObjectStorage::TEvInferredFileSchema, TMetadataResult>(
374375
arrowInferencinatorId,
375-
new NObjectStorage::TEvInferFileSchema(TString{pathFut.GetValue()}),
376+
new NObjectStorage::TEvInferFileSchema(TString{path}, size),
376377
promise,
377378
std::move(schemaToMetadata)
378379
));

‎ydb/core/external_sources/object_storage/events.h‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ enum EEventTypes : ui32 {
2020
EvFileError,
2121

2222
EvRequestS3Range,
23+
EvRequestS3Schema,
2324
EvS3DownloadResponse,
2425
EvS3RangeResponse,
2526
EvS3RangeError,
@@ -29,6 +30,7 @@ enum EEventTypes : ui32 {
2930
EvInferredFileSchema,
3031

3132
EvArrowFile,
33+
EvArrowSchema,
3234

3335
EvEnd,
3436
};
@@ -68,6 +70,22 @@ struct TEvRequestS3Range : public NActors::TEventLocal<TEvRequestS3Range, EvRequ
6870
NActors::TActorId Sender;
6971
};
7072

73+
struct TEvRequestS3Schema : public NActors::TEventLocal<TEvRequestS3Schema, EvRequestS3Schema> {
74+
TEvRequestS3Schema(TString path, ui64 size, const TGUID& requestId, NActors::TActorId sender)
75+
: Path{std::move(path)}
76+
, Size{size}
77+
, RequestId{requestId}
78+
, Sender{sender}
79+
{}
80+
81+
TString Path;
82+
83+
ui64 Size;
84+
85+
TGUID RequestId;
86+
NActors::TActorId Sender;
87+
};
88+
7189
struct TEvS3DownloadResponse : public NActors::TEventLocal<TEvS3DownloadResponse, EvS3DownloadResponse> {
7290
TEvS3DownloadResponse(std::shared_ptr<TEvRequestS3Range>&& request, NYql::IHTTPGateway::TResult&& result)
7391
: Request{std::move(request)}
@@ -118,12 +136,24 @@ struct TEvArrowFile : public NActors::TEventLocal<TEvArrowFile, EvArrowFile> {
118136
TString Path;
119137
};
120138

139+
struct TEvArrowSchema : public NActors::TEventLocal<TEvArrowSchema, EvArrowSchema> {
140+
TEvArrowSchema(std::shared_ptr<arrow::Schema> schema, TString path)
141+
: Schema{std::move(schema)}
142+
, Path{std::move(path)}
143+
{}
144+
145+
std::shared_ptr<arrow::Schema> Schema;
146+
TString Path;
147+
};
148+
121149
struct TEvInferFileSchema : public NActors::TEventLocal<TEvInferFileSchema, EvInferFileSchema> {
122-
explicit TEvInferFileSchema(TString&& path)
150+
explicit TEvInferFileSchema(TString&& path, ui64 size)
123151
: Path{std::move(path)}
152+
, Size{size}
124153
{}
125154

126155
TString Path;
156+
ui64 Size = 0;
127157
};
128158

129159
struct TEvInferredFileSchema : public NActors::TEventLocal<TEvInferredFileSchema, EvInferredFileSchema> {

‎ydb/core/external_sources/object_storage/inference/arrow_inferencinator.cpp‎

Lines changed: 57 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -240,10 +240,15 @@ std::unique_ptr<FormatConfig> MakeFormatConfig(EFileFormat format, const THashMa
240240

241241
class TArrowInferencinator : public NActors::TActorBootstrapped<TArrowInferencinator> {
242242
public:
243-
TArrowInferencinator(NActors::TActorId arrowFetcher, EFileFormat format, const THashMap<TString, TString>& params)
243+
TArrowInferencinator(
244+
NActors::TActorId arrowFetcher,
245+
NActors::TActorId s3Fetcher,
246+
EFileFormat format,
247+
const THashMap<TString, TString>& params)
244248
: Format_{format}
245249
, Config_{MakeFormatConfig(Format_, params)}
246250
, ArrowFetcherId_{arrowFetcher}
251+
, S3FetcherId_{s3Fetcher}
247252
{
248253
Y_ABORT_UNLESS(IsArrowInferredFormat(Format_));
249254
}
@@ -256,11 +261,34 @@ class TArrowInferencinator : public NActors::TActorBootstrapped<TArrowInferencin
256261
HFunc(TEvInferFileSchema, HandleInferRequest);
257262
HFunc(TEvFileError, HandleFileError);
258263
HFunc(TEvArrowFile, HandleFileInference);
264+
HFunc(TEvArrowSchema, HandleFileSchema);
259265
)
260266

261267
void HandleInferRequest(TEvInferFileSchema::TPtr& ev, const NActors::TActorContext& ctx) {
262268
RequesterId_ = ev->Sender;
263-
ctx.Send(ArrowFetcherId_, ev->Release());
269+
auto& event = *ev->Get();
270+
271+
switch (Format_) {
272+
case EFileFormat::CsvWithNames:
273+
case EFileFormat::TsvWithNames: {
274+
ctx.Send(ArrowFetcherId_, ev->Release());
275+
return;
276+
}
277+
case EFileFormat::Parquet: {
278+
auto s3SchemaRequest = new TEvRequestS3Schema(event.Path, event.Size, {}, SelfId());
279+
CreateGuid(&s3SchemaRequest->RequestId);
280+
281+
ctx.Send(S3FetcherId_, s3SchemaRequest);
282+
return;
283+
}
284+
default: {
285+
ctx.Send(RequesterId_, MakeError(event.Path, NFq::TIssuesIds::UNSUPPORTED, TStringBuilder{} << "unsupported format for inference: " << ConvertFileFormat(Format_)));
286+
return;
287+
}
288+
case EFileFormat::Undefined:
289+
Y_ABORT("Invalid format should be unreachable");
290+
}
291+
264292
}
265293

266294
void HandleFileInference(TEvArrowFile::TPtr& ev, const NActors::TActorContext& ctx) {
@@ -271,33 +299,48 @@ class TArrowInferencinator : public NActors::TActorBootstrapped<TArrowInferencin
271299
return;
272300
}
273301

274-
auto& arrowFields = std::get<ArrowFields>(mbArrowFields);
302+
ConvertArrowSchema(std::get<ArrowFields>(mbArrowFields), file.Path, ctx);
303+
}
304+
305+
void HandleFileSchema(TEvArrowSchema::TPtr& ev, const NActors::TActorContext& ctx) {
306+
auto& schema = *ev->Get();
307+
308+
ConvertArrowSchema(schema.Schema->fields(), schema.Path, ctx);
309+
}
310+
311+
void HandleFileError(TEvFileError::TPtr& ev, const NActors::TActorContext& ctx) {
312+
Cout << "TArrowInferencinator::HandleFileError" << Endl;
313+
ctx.Send(RequesterId_, new TEvInferredFileSchema(ev->Get()->Path, std::move(ev->Get()->Issues)));
314+
}
315+
316+
private:
317+
void ConvertArrowSchema(const ArrowFields& fields, const TString& path, const NActors::TActorContext& ctx) const {
275318
std::vector<Ydb::Column> ydbFields;
276-
for (const auto& field : arrowFields) {
319+
for (const auto& field : fields) {
277320
ydbFields.emplace_back();
278321
auto& ydbField = ydbFields.back();
279322
if (!ArrowToYdbType(*ydbField.mutable_type(), *field->type())) {
280-
ctx.Send(RequesterId_, MakeErrorSchema(file.Path, NFq::TIssuesIds::UNSUPPORTED, TStringBuilder{} << "couldn't convert arrow type to ydb: " << field->ToString()));
323+
ctx.Send(RequesterId_, MakeErrorSchema(path, NFq::TIssuesIds::UNSUPPORTED, TStringBuilder{} << "couldn't convert arrow type to ydb: " << field->ToString()));
281324
return;
282325
}
283326
ydbField.mutable_name()->assign(field->name());
284327
}
285-
ctx.Send(RequesterId_, new TEvInferredFileSchema(file.Path, std::move(ydbFields)));
328+
ctx.Send(RequesterId_, new TEvInferredFileSchema(path, std::move(ydbFields)));
286329
}
287330

288-
void HandleFileError(TEvFileError::TPtr& ev, const NActors::TActorContext& ctx) {
289-
Cout << "TArrowInferencinator::HandleFileError" << Endl;
290-
ctx.Send(RequesterId_, new TEvInferredFileSchema(ev->Get()->Path, std::move(ev->Get()->Issues)));
291-
}
292-
293-
private:
294331
EFileFormat Format_;
295332
std::unique_ptr<FormatConfig> Config_;
296333
NActors::TActorId ArrowFetcherId_;
334+
NActors::TActorId S3FetcherId_;
297335
NActors::TActorId RequesterId_;
298336
};
299337

300-
NActors::IActor* CreateArrowInferencinator(NActors::TActorId arrowFetcher, EFileFormat format, const THashMap<TString, TString>& params) {
301-
return new TArrowInferencinator{arrowFetcher, format, params};
338+
NActors::IActor* CreateArrowInferencinator(
339+
NActors::TActorId arrowFetcher,
340+
NActors::TActorId s3Fetcher,
341+
EFileFormat format,
342+
const THashMap<TString, TString>& params) {
343+
344+
return new TArrowInferencinator{arrowFetcher, s3Fetcher, format, params};
302345
}
303346
} // namespace NKikimr::NExternalSource::NObjectStorage::NInference

‎ydb/core/external_sources/object_storage/inference/arrow_inferencinator.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,5 +53,5 @@ constexpr bool IsArrowInferredFormat(TStringBuf format) {
5353
return IsArrowInferredFormat(ConvertFileFormat(format));
5454
}
5555

56-
NActors::IActor* CreateArrowInferencinator(NActors::TActorId arrowFetcher, EFileFormat format, const THashMap<TString, TString>& params);
56+
NActors::IActor* CreateArrowInferencinator(NActors::TActorId arrowFetcher, NActors::TActorId s3Fetcher, EFileFormat format, const THashMap<TString, TString>& params);
5757
} // namespace NKikimr::NExternalSource::NObjectStorage::NInference

‎ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class ArrowInferenceTest : public testing::Test {
5151
NActors::TActorId RegisterInferencinator(TStringBuf formatStr) {
5252
auto format = NInference::ConvertFileFormat(formatStr);
5353
auto arrowFetcher = ActorSystem.Register(NInference::CreateArrowFetchingActor(S3ActorId, format, {}), 1);
54-
return ActorSystem.Register(NInference::CreateArrowInferencinator(arrowFetcher, format, {}), 1);
54+
return ActorSystem.Register(NInference::CreateArrowInferencinator(arrowFetcher, S3ActorId, format, {}), 1);
5555
}
5656

5757
void TearDown() override {
@@ -85,7 +85,7 @@ TEST_F(ArrowInferenceTest, csv_simple) {
8585

8686
auto inferencinatorId = RegisterInferencinator("csv_with_names");
8787
ActorSystem.WrapInActorContext(EdgeActorId, [this, inferencinatorId] {
88-
NActors::TActivationContext::AsActorContext().Send(inferencinatorId, new TEvInferFileSchema(TString{Path}));
88+
NActors::TActivationContext::AsActorContext().Send(inferencinatorId, new TEvInferFileSchema(TString{Path}, 0));
8989
});
9090

9191
std::unique_ptr<NActors::IEventHandle> event = ActorSystem.WaitForEdgeActorEvent({EdgeActorId});
@@ -121,7 +121,7 @@ TEST_F(ArrowInferenceTest, tsv_simple) {
121121

122122
auto inferencinatorId = RegisterInferencinator("tsv_with_names");
123123
ActorSystem.WrapInActorContext(EdgeActorId, [this, inferencinatorId] {
124-
NActors::TActivationContext::AsActorContext().Send(inferencinatorId, new TEvInferFileSchema(TString{Path}));
124+
NActors::TActivationContext::AsActorContext().Send(inferencinatorId, new TEvInferFileSchema(TString{Path}, 0));
125125
});
126126

127127
std::unique_ptr<NActors::IEventHandle> event = ActorSystem.WaitForEdgeActorEvent({EdgeActorId});

‎ydb/core/external_sources/object_storage/inference/ya.make‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ PEERDIR(
1717
ydb/core/external_sources/object_storage
1818

1919
ydb/library/yql/providers/s3/compressors
20+
ydb/library/yql/providers/common/arrow
2021
)
2122

2223
END()

‎ydb/core/external_sources/object_storage/s3_fetcher.cpp‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "s3_fetcher.h"
22

33
#include <ydb/library/actors/core/hfunc.h>
4+
#include <ydb/library/yql/providers/common/arrow/interface/arrow_reader.h>
45

56
namespace NKikimr::NExternalSource::NObjectStorage {
67

@@ -23,6 +24,7 @@ class S3Fetcher : public NActors::TActorBootstrapped<S3Fetcher> {
2324

2425
STRICT_STFUNC(WorkingState,
2526
HFunc(TEvRequestS3Range, HandleRequest);
27+
HFunc(TEvRequestS3Schema, HandleRequest);
2628

2729
HFunc(TEvS3DownloadResponse, HandleDownloadReponse);
2830
)
@@ -31,6 +33,10 @@ class S3Fetcher : public NActors::TActorBootstrapped<S3Fetcher> {
3133
StartDownload(std::shared_ptr<TEvRequestS3Range>(ev->Release().Release()), ctx.ActorSystem());
3234
}
3335

36+
void HandleRequest(TEvRequestS3Schema::TPtr& ev, const NActors::TActorContext& ctx) {
37+
StartDownload(std::shared_ptr<TEvRequestS3Schema>(ev->Release().Release()), ctx.ActorSystem());
38+
}
39+
3440
void HandleDownloadReponse(TEvS3DownloadResponse::TPtr& ev, const NActors::TActorContext& ctx) {
3541
auto& response = *ev->Get();
3642
auto& result = response.Result;
@@ -76,6 +82,32 @@ class S3Fetcher : public NActors::TActorBootstrapped<S3Fetcher> {
7682
}, {}, RetryPolicy_);
7783
}
7884

85+
void StartDownload(std::shared_ptr<TEvRequestS3Schema>&& request, NActors::TActorSystem* actorSystem) {
86+
const auto& authInfo = Credentials_.GetAuthInfo();
87+
auto headers = NYql::IHTTPGateway::MakeYcHeaders(
88+
request->RequestId.AsGuidString(),
89+
authInfo.GetToken(),
90+
{},
91+
authInfo.GetAwsUserPwd(),
92+
authInfo.GetAwsSigV4()
93+
);
94+
95+
NYql::TArrowFileDesc desc(
96+
Url_ + request->Path,
97+
Gateway_,
98+
std::move(headers),
99+
RetryPolicy_,
100+
request->Size,
101+
"parquet"
102+
);
103+
104+
auto schemaReader = NYql::MakeArrowReader(NYql::TArrowReaderSettings());
105+
auto futureSchema = schemaReader->GetSchema(desc);
106+
futureSchema.Apply([actorSystem, request](NThreading::TFuture<NYql::IArrowReader::TSchemaResponse> response) {
107+
actorSystem->Send(request->Sender, new TEvArrowSchema(response.GetValue().Schema, request->Path));
108+
});
109+
}
110+
79111
private:
80112
TString Url_;
81113
NYql::IHTTPGateway::TPtr Gateway_;

0 commit comments

Comments
 (0)