Merge 55276be into 2848978

Author: evanevanevanevannnn

ydb/core/external_sources/object_storage.cpp

@@ -3,6 +3,7 @@
 #include "validation_functions.h"
 #include "object_storage/s3_fetcher.h"
 
+#include <util/string/join.h>
 #include <ydb/core/external_sources/object_storage/inference/arrow_fetcher.h>
 #include <ydb/core/external_sources/object_storage/inference/arrow_inferencinator.h>
 #include <ydb/core/kqp/gateway/actors/kqp_ic_gateway_actors.h>
@@ -20,6 +21,10 @@
 #include <ydb/public/sdk/cpp/client/ydb_value/value.h>
 
 #include <library/cpp/scheme/scheme.h>
+#include <library/cpp/json/json_reader.h>
+#include <arrow/buffer_builder.h>
+#include <arrow/buffer.h>
+#include <arrow/io/memory.h>
 
 #include <util/string/builder.h>
 
@@ -322,21 +327,26 @@ struct TObjectStorageExternalSource : public IExternalSource {
             structuredTokenBuilder.SetNoAuth();
         }
 
-        auto effectiveFilePattern = NYql::NS3::NormalizePath(meta->TableLocation);
-        if (meta->TableLocation.EndsWith('/')) {
-            effectiveFilePattern += '*';
-        } 
-
         const NYql::TS3Credentials credentials(CredentialsFactory, structuredTokenBuilder.ToJson());
+
+        const TString path = meta->TableLocation;
+        const TString filePattern = meta->Attributes.Value("filepattern", TString{});
+        const TVector<TString> partitionedBy = GetPartitionedByConfig(meta);
+        NYql::NS3Lister::TListingRequest request {
+            .Url = meta->DataSourceLocation,
+            .Credentials = credentials
+        };
+
+        NYql::NS3::BuildS3FilePattern(path, filePattern, partitionedBy, request);
+
+        auto partByData = std::make_shared<TStringBuilder>();
+
         auto httpGateway = NYql::IHTTPGateway::Make();
         auto httpRetryPolicy = NYql::GetHTTPDefaultRetryPolicy(NYql::THttpRetryPolicyOptions{.RetriedCurlCodes = NYql::FqRetriedCurlCodes()});
-        auto s3Lister = NYql::NS3Lister::MakeS3Lister(httpGateway, httpRetryPolicy, NYql::NS3Lister::TListingRequest{
-            .Url = meta->DataSourceLocation,
-            .Credentials = credentials,
-            .Pattern = effectiveFilePattern,
-        }, Nothing(), AllowLocalFiles, ActorSystem);
-        auto afterListing = s3Lister->Next().Apply([path = effectiveFilePattern](const NThreading::TFuture<NYql::NS3Lister::TListResult>& listResFut) {
+        auto s3Lister = NYql::NS3Lister::MakeS3Lister(httpGateway, httpRetryPolicy, request, Nothing(), AllowLocalFiles, ActorSystem);
+        auto afterListing = s3Lister->Next().Apply([partByData, partitionedBy, path = request.Pattern](const NThreading::TFuture<NYql::NS3Lister::TListResult>& listResFut) {
             auto& listRes = listResFut.GetValue();
+            auto& partByRef = *partByData;
             if (std::holds_alternative<NYql::NS3Lister::TListError>(listRes)) {
                 auto& error = std::get<NYql::NS3Lister::TListError>(listRes);
                 throw yexception() << error.Issues.ToString();
@@ -345,6 +355,14 @@ struct TObjectStorageExternalSource : public IExternalSource {
             if (entries.Objects.empty()) {
                 throw yexception() << "couldn't find files at " << path;
             }
+
+            partByRef << JoinSeq(",", partitionedBy);
+            for (const auto& entry : entries.Objects) {
+                if (entry.MatchedGlobs.size() != partitionedBy.size()) {
+                    continue;
+                }
+                partByRef << Endl << JoinSeq(",", entry.MatchedGlobs);
+            }
             for (const auto& entry : entries.Objects) {
                 if (entry.Size > 0) {
                     return entry;
@@ -387,12 +405,14 @@ struct TObjectStorageExternalSource : public IExternalSource {
             auto [path, size, _] = entryFut.GetValue();
             actorSystem->Register(new NKqp::TActorRequestHandler<NObjectStorage::TEvInferFileSchema, NObjectStorage::TEvInferredFileSchema, TMetadataResult>(
                 arrowInferencinatorId,
-                new NObjectStorage::TEvInferFileSchema(TString{path}, size),
+                new NObjectStorage::TEvInferFileSchema(std::move(path), size),
                 promise,
                 std::move(schemaToMetadata)
             ));
 
             return promise.GetFuture();
+        }).Apply([arrowInferencinatorId, meta, partByData, partitionedBy, this](const NThreading::TFuture<TMetadataResult>& result) {
+            return InferPartitionedColumnsTypes(arrowInferencinatorId, partByData, partitionedBy, result);
         }).Apply([](const NThreading::TFuture<TMetadataResult>& result) {
             auto& value = result.GetValue();
             if (value.Success()) {
@@ -407,6 +427,85 @@ struct TObjectStorageExternalSource : public IExternalSource {
     }
 
 private:
+    NThreading::TFuture<TMetadataResult> InferPartitionedColumnsTypes(
+        NActors::TActorId arrowInferencinatorId,
+        std::shared_ptr<TStringBuilder> partByData,
+        const TVector<TString>& partitionedBy,
+        const NThreading::TFuture<TMetadataResult>& result) const {
+
+        auto& value = result.GetValue();
+        if (!value.Success() || partitionedBy.empty()) {
+            return result;
+        }
+
+        auto meta = value.Metadata;
+        for (const auto& partitionName : partitionedBy) {
+            auto& destColumn = *meta->Schema.add_column();
+            destColumn.mutable_name()->assign(partitionName);
+            destColumn.mutable_type()->set_type_id(Ydb::Type::UTF8);
+        }
+
+        arrow::BufferBuilder builder;
+        auto partitionBuffer = std::make_shared<arrow::Buffer>(nullptr, 0);
+        auto buildStatus = builder.Append(partByData->data(), partByData->size());
+        auto finishStatus = builder.Finish(&partitionBuffer);
+
+        if (!buildStatus.ok() || !finishStatus.ok()) {
+            return result;
+        }
+
+        auto promise = NThreading::NewPromise<TMetadataResult>();
+        auto partitionsToMetadata = [meta](NThreading::TPromise<TMetadataResult> metaPromise, NObjectStorage::TEvInferredFileSchema&& response){
+            if (response.Status.IsSuccess()) {
+                THashMap<TString, Ydb::Type> inferredTypes;
+                for (const auto& column : response.Fields) {
+                    inferredTypes[column.name()] = column.type();
+                }
+                
+                for (auto& destColumn : *meta->Schema.mutable_column()) {
+                    if (auto type = inferredTypes.FindPtr(destColumn.name()); type) {
+                        destColumn.mutable_type()->set_type_id(type->has_optional_type() ? 
+                            type->optional_type().item().type_id() : 
+                            type->type_id());
+                    }
+                }
+            }
+            TMetadataResult result;
+            result.SetSuccess();
+            result.Metadata = meta;
+            metaPromise.SetValue(std::move(result));
+        };
+
+        auto file = std::make_shared<arrow::io::BufferReader>(std::move(partitionBuffer));
+        ActorSystem->Register(new NKqp::TActorRequestHandler<NObjectStorage::TEvInferPartitions, NObjectStorage::TEvInferredFileSchema, TMetadataResult>(
+            arrowInferencinatorId,
+            new NObjectStorage::TEvInferPartitions(std::move(std::dynamic_pointer_cast<arrow::io::RandomAccessFile>(file))),
+            promise,
+            std::move(partitionsToMetadata)
+        ));
+
+        return promise.GetFuture();
+    }
+
+    static TVector<TString> GetPartitionedByConfig(std::shared_ptr<TMetadata> meta) {
+        THashSet<TString> columns;
+        if (auto partitioned = meta->Attributes.FindPtr("partitionedby"); partitioned) {
+            NJson::TJsonValue values;
+            Y_ENSURE(NJson::ReadJsonTree(*partitioned, &values));
+            Y_ENSURE(values.GetType() == NJson::JSON_ARRAY);
+
+            for (const auto& value : values.GetArray()) {
+                Y_ENSURE(value.GetType() == NJson::JSON_STRING);
+                if (columns.find(value.GetString()) != columns.end()) {
+                    throw yexception() << "invalid partitioned_by parameter, column " << value.GetString() << "mentioned twice";
+                }
+                columns.insert(value.GetString());
+            }
+        }
+
+        return TVector<TString>{columns.begin(), columns.end()};
+    }
+
     static bool IsValidIntervalUnit(const TString& unit) {
         static constexpr std::array<std::string_view, 7> IntervalUnits = {
             "MICROSECONDS"sv,

ydb/core/external_sources/object_storage/events.h

@@ -26,6 +26,7 @@ enum EEventTypes : ui32 {
 
     // Move somewhere separate?
     EvInferFileSchema,
+    EvInferPartitions,
     EvInferredFileSchema,
 
     EvArrowFile,
@@ -128,6 +129,14 @@ struct TEvInferFileSchema : public NActors::TEventLocal<TEvInferFileSchema, EvIn
     ui64 Size = 0;
 };
 
+struct TEvInferPartitions : public NActors::TEventLocal<TEvInferPartitions, EvInferPartitions> {
+    explicit TEvInferPartitions(std::shared_ptr<arrow::io::RandomAccessFile> file)
+        : File{std::move(file)}
+    {}
+
+    std::shared_ptr<arrow::io::RandomAccessFile> File;
+};
+
 struct TEvInferredFileSchema : public NActors::TEventLocal<TEvInferredFileSchema, EvInferredFileSchema> {
     TEvInferredFileSchema(TString path, std::vector<Ydb::Column>&& fields)
         : Path{std::move(path)}

ydb/core/external_sources/object_storage/inference/arrow_fetcher.cpp

@@ -1,5 +1,6 @@
 #include "arrow_fetcher.h"
 #include "arrow_inferencinator.h"
+#include "infer_config.h"
 
 #include <arrow/buffer.h>
 #include <arrow/buffer_builder.h>
@@ -30,6 +31,7 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
     TArrowFileFetcher(NActors::TActorId s3FetcherId, EFileFormat format, const THashMap<TString, TString>& params)
         : S3FetcherId_{s3FetcherId}
         , Format_{format}
+        , Config_{MakeFormatConfig(format, params)}
     {
         Y_ABORT_UNLESS(IsArrowInferredFormat(Format_));
 
@@ -72,7 +74,11 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
                 break;
             }
             default: {
-                ctx.Send(localRequest.Requester, MakeError(localRequest.Path, NFq::TIssuesIds::UNSUPPORTED, TStringBuilder{} << "unsupported format for inference: " << ConvertFileFormat(Format_)));
+                ctx.Send(localRequest.Requester, MakeError(
+                    localRequest.Path,
+                    NFq::TIssuesIds::UNSUPPORTED,
+                    TStringBuilder{} << "unsupported format for inference: " << ConvertFileFormat(Format_))
+                );
                 return;
             }
             case EFileFormat::Undefined:
@@ -100,12 +106,7 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
         switch (Format_) {
             case EFileFormat::CsvWithNames:
             case EFileFormat::TsvWithNames: {
-                // TODO: obtain from request
-                arrow::csv::ParseOptions options;
-                if (Format_ == EFileFormat::TsvWithNames) {
-                    options.delimiter = '\t';
-                }
-                file = CleanupCsvFile(data, request, options, ctx);
+                file = CleanupCsvFile(data, request, std::dynamic_pointer_cast<CsvConfig>(Config_)->ParseOpts, ctx);
                 ctx.Send(request.Requester, new TEvArrowFile(std::move(file), request.Path));
                 break;
             }
@@ -120,7 +121,7 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
             }
             case EFileFormat::JsonEachRow:
             case EFileFormat::JsonList: {
-                file = CleanupJsonFile(data, request, arrow::json::ParseOptions::Defaults(), ctx);
+                file = CleanupJsonFile(data, request, std::dynamic_pointer_cast<JsonConfig>(Config_)->ParseOpts, ctx);
                 ctx.Send(request.Requester, new TEvArrowFile(std::move(file), request.Path));
                 break;
             }
@@ -354,6 +355,7 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
     // Fields
     NActors::TActorId S3FetcherId_;
     EFileFormat Format_;
+    std::shared_ptr<FormatConfig> Config_;
     TMaybe<TString> DecompressionFormat_;
     std::unordered_map<TString, TRequest> InflightRequests_; // Path -> Request
 };