Merge 947bc67 into 5fe6486

alchizhevsky · web-flow · commit 1b98ee8426d8 · 2024-04-26T11:01:06.000Z
diff --git a/ydb/library/yql/providers/s3/actors/yql_s3_read_actor.cpp b/ydb/library/yql/providers/s3/actors/yql_s3_read_actor.cpp
@@ -160,20 +160,20 @@ struct TEvS3FileQueue {
 
         EvEnd
     };
-    static_assert(EvEnd < EventSpaceEnd(NKikimr::TKikimrEvents::ES_S3_FILE_QUEUE), 
+    static_assert(EvEnd < EventSpaceEnd(NKikimr::TKikimrEvents::ES_S3_FILE_QUEUE),
                   "expect EvEnd < EventSpaceEnd(TEvents::ES_S3_FILE_QUEUE)");
-    
+
     struct TEvUpdateConsumersCount :
         public TEventPB<TEvUpdateConsumersCount, NS3::FileQueue::TEvUpdateConsumersCount, EvUpdateConsumersCount> {
-        
+
         explicit TEvUpdateConsumersCount(ui64 consumersCountDelta = 0) {
             Record.SetConsumersCountDelta(consumersCountDelta);
         }
     };
 
     struct TEvAck :
         public TEventPB<TEvAck, NS3::FileQueue::TEvAck, EvAck> {
-        
+
         TEvAck() = default;
 
         explicit TEvAck(const TMessageTransportMeta& transportMeta) {
@@ -388,6 +388,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
         TPathList paths,
         size_t prefetchSize,
         ui64 fileSizeLimit,
+        ui64 readLimit,
         bool useRuntimeListing,
         ui64 consumersCount,
         ui64 batchSizeLimit,
@@ -401,6 +402,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
         : TxId(std::move(txId))
         , PrefetchSize(prefetchSize)
         , FileSizeLimit(fileSizeLimit)
+        , ReadLimit(readLimit)
         , MaybeIssues(Nothing())
         , UseRuntimeListing(useRuntimeListing)
         , ConsumersCount(consumersCount)
@@ -510,7 +512,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
                 // skip 'directories'
                 continue;
             }
-            if (object.Size > FileSizeLimit) {
+            if (object.Size > std::min(FileSizeLimit, ReadLimit)) {
                 auto errorMessage = TStringBuilder()
                                     << "Size of object " << object.Path << " = "
                                     << object.Size
@@ -522,10 +524,10 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
             LOG_T("TS3FileQueueActor", "SaveRetrievedResults adding path: " << object.Path);
             TObjectPath objectPath;
             objectPath.SetPath(object.Path);
-            objectPath.SetSize(object.Size);
+            objectPath.SetSize(std::min(object.Size, ReadLimit));
             objectPath.SetPathIndex(CurrentDirectoryPathIndex);
             Objects.emplace_back(std::move(objectPath));
-            ObjectsTotalSize += object.Size;
+            ObjectsTotalSize += std::min(object.Size, ReadLimit);
         }
         return true;
     }
@@ -595,7 +597,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
         Send(ev->Sender, new TEvS3FileQueue::TEvObjectPathReadError(*MaybeIssues, ev->Get()->Record.GetTransportMeta()));
         TryFinish(ev->Sender, ev->Get()->Record.GetTransportMeta().GetSeqNo());
     }
-    
+
     void HandleUpdateConsumersCount(TEvS3FileQueue::TEvUpdateConsumersCount::TPtr& ev) {
         if (!UpdatedConsumers.contains(ev->Sender)) {
             LOG_D(
@@ -649,7 +651,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
 
         LOG_T("TS3FileQueueActor", "SendObjects Sending " << result.size() << " objects to consumer with id " << consumer);
         Send(consumer, new TEvS3FileQueue::TEvObjectPathBatch(std::move(result), HasNoMoreItems(), transportMeta));
-        
+
         if (HasNoMoreItems()) {
             TryFinish(consumer, transportMeta.GetSeqNo());
         }
@@ -671,7 +673,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
     }
 
     bool CanSendToConsumer(const TActorId& consumer) {
-        return !UseRuntimeListing || RoundRobinStageFinished || 
+        return !UseRuntimeListing || RoundRobinStageFinished ||
                (StartedConsumers.size() < ConsumersCount && !StartedConsumers.contains(consumer));
     }
 
@@ -749,7 +751,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
                     }
                 });
     }
-    
+
     void ScheduleRequest(const TActorId& consumer, const TMessageTransportMeta& transportMeta) {
         PendingRequests[consumer].push_back(transportMeta);
         HasPendingRequests = true;
@@ -786,7 +788,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
             }
         }
     }
-    
+
     void TryFinish(const TActorId& consumer, ui64 seqNo) {
         LOG_T("TS3FileQueueActor", "TryFinish from consumer " << consumer << ", " << FinishedConsumers.size() << " consumers already finished, seqNo=" << seqNo);
         if (FinishingConsumerToLastSeqNo.contains(consumer)) {
@@ -810,6 +812,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
 
     size_t PrefetchSize;
     ui64 FileSizeLimit;
+    ui64 ReadLimit;
     TMaybe<NS3Lister::IS3Lister::TPtr> MaybeLister = Nothing();
     TMaybe<NThreading::TFuture<NS3Lister::TListResult>> ListingFuture;
     size_t CurrentDirectoryPathIndex = 0;
@@ -834,7 +837,7 @@ class TS3FileQueueActor : public TActorBootstrapped<TS3FileQueueActor> {
     const TString Pattern;
     const ES3PatternVariant PatternVariant;
     const ES3PatternType PatternType;
-    
+
     static constexpr TDuration PoisonTimeout = TDuration::Hours(3);
     static constexpr TDuration RoundRobinStageTimeout = TDuration::Seconds(3);
 };
@@ -914,6 +917,7 @@ class TS3ReadActor : public TActorBootstrapped<TS3ReadActor>, public IDqComputeA
                 std::move(Paths),
                 ReadActorFactoryCfg.MaxInflight * 2,
                 FileSizeLimit,
+                SizeLimit,
                 false,
                 1,
                 FileQueueBatchSizeLimit,
@@ -1093,7 +1097,7 @@ class TS3ReadActor : public TActorBootstrapped<TS3ReadActor>, public IDqComputeA
     void HandleAck(TEvS3FileQueue::TEvAck::TPtr& ev) {
         FileQueueEvents.OnEventReceived(ev);
     }
-    
+
     static void OnDownloadFinished(TActorSystem* actorSystem, TActorId selfId, const TString& requestId, IHTTPGateway::TResult&& result, size_t pathInd, const TString path) {
         if (!result.Issues) {
             actorSystem->Send(new IEventHandle(selfId, TActorId(), new TEvPrivate::TEvReadResult(std::move(result.Content), requestId, pathInd, path)));
@@ -1205,7 +1209,7 @@ class TS3ReadActor : public TActorBootstrapped<TS3ReadActor>, public IDqComputeA
         auto issues = NS3Util::AddParentIssue(TStringBuilder{} << "Error while reading file " << path << " with request id [" << requestId << "]", TIssues{result->Get()->Error});
         Send(ComputeActorId, new TEvAsyncInputError(InputIndex, std::move(issues), NYql::NDqProto::StatusIds::EXTERNAL_ERROR));
     }
-    
+
     void Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvRetry::TPtr&) {
         FileQueueEvents.Retry();
     }
@@ -2084,7 +2088,7 @@ class TS3ReadCoroImpl : public TActorCoroImpl {
                 if (isCancelled) {
                     LOG_CORO_D("RunCoroBlockArrowParserOverHttp - STOPPED ON SATURATION, downloaded " <<
                                QueueBufferCounter->DownloadedBytes << " bytes");
-                    break;        
+                    break;
                 }
             }
         }
@@ -2534,6 +2538,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
         ::NMonitoring::TDynamicCounterPtr counters,
         ::NMonitoring::TDynamicCounterPtr taskCounters,
         ui64 fileSizeLimit,
+        ui64 readLimit,
         std::optional<ui64> rowsLimitHint,
         IMemoryQuotaManager::TPtr memoryQuotaManager,
         bool useRuntimeListing,
@@ -2560,6 +2565,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
         , TaskCounters(std::move(taskCounters))
         , FileQueueActor(fileQueueActor)
         , FileSizeLimit(fileSizeLimit)
+        , ReadLimit(readLimit)
         , MemoryQuotaManager(memoryQuotaManager)
         , UseRuntimeListing(useRuntimeListing)
         , FileQueueBatchSizeLimit(fileQueueBatchSizeLimit)
@@ -2618,6 +2624,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
                 std::move(Paths),
                 ReadActorFactoryCfg.MaxInflight * 2,
                 FileSizeLimit,
+                ReadLimit,
                 false,
                 1,
                 FileQueueBatchSizeLimit,
@@ -2780,7 +2787,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
     void CommitState(const NDqProto::TCheckpoint&) final {}
 
     ui64 GetInputIndex() const final {
-        return InputIndex; 
+        return InputIndex;
     }
 
     const TDqAsyncStats& GetIngressStats() const final {
@@ -3034,7 +3041,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
             }
         }
     }
-    
+
     void Handle(TEvS3FileQueue::TEvAck::TPtr& ev) {
         FileQueueEvents.OnEventReceived(ev);
     }
@@ -3132,6 +3139,7 @@ class TS3StreamReadActor : public TActorBootstrapped<TS3StreamReadActor>, public
     std::set<NActors::TActorId> CoroActors;
     NActors::TActorId FileQueueActor;
     const ui64 FileSizeLimit;
+    const ui64 ReadLimit;
     bool Bootstrapped = false;
     IMemoryQuotaManager::TPtr MemoryQuotaManager;
     bool UseRuntimeListing;
@@ -3291,6 +3299,7 @@ IActor* CreateS3FileQueueActor(
         TPathList paths,
         size_t prefetchSize,
         ui64 fileSizeLimit,
+        ui64 readLimit,
         bool useRuntimeListing,
         ui64 consumersCount,
         ui64 batchSizeLimit,
@@ -3306,6 +3315,7 @@ IActor* CreateS3FileQueueActor(
         paths,
         prefetchSize,
         fileSizeLimit,
+        readLimit,
         useRuntimeListing,
         consumersCount,
         batchSizeLimit,
@@ -3390,15 +3400,15 @@ std::pair<NYql::NDq::IDqComputeActorAsyncInput*, IActor*> CreateS3ReadActor(
     if (params.GetRowsLimitHint() != 0) {
         rowsLimitHint = params.GetRowsLimitHint();
     }
-    
+
     TActorId fileQueueActor;
     if (auto it = settings.find("fileQueueActor"); it != settings.cend()) {
         NActorsProto::TActorId protoId;
         TMemoryInput inputStream(it->second);
         ParseFromTextFormat(inputStream, protoId);
         fileQueueActor = ActorIdFromProto(protoId);
     }
-    
+
     ui64 fileQueueBatchSizeLimit = 0;
     if (auto it = settings.find("fileQueueBatchSizeLimit"); it != settings.cend()) {
         fileQueueBatchSizeLimit = FromString<ui64>(it->second);
@@ -3408,7 +3418,7 @@ std::pair<NYql::NDq::IDqComputeActorAsyncInput*, IActor*> CreateS3ReadActor(
     if (auto it = settings.find("fileQueueBatchObjectCountLimit"); it != settings.cend()) {
         fileQueueBatchObjectCountLimit = FromString<ui64>(it->second);
     }
-    
+
     ui64 fileQueueConsumersCountDelta = 0;
     if (readRanges.size() > 1) {
         fileQueueConsumersCountDelta = readRanges.size() - 1;
@@ -3516,9 +3526,13 @@ std::pair<NYql::NDq::IDqComputeActorAsyncInput*, IActor*> CreateS3ReadActor(
 
 #undef SET_FLAG
 #undef SUPPORTED_FLAGS
+        ui64 sizeLimit = std::numeric_limits<ui64>::max();
+        if (const auto it = settings.find("sizeLimit"); settings.cend() != it)
+            sizeLimit = FromString<ui64>(it->second);
+
         const auto actor = new TS3StreamReadActor(inputIndex, statsLevel, txId, std::move(gateway), holderFactory, params.GetUrl(), authInfo, pathPattern, pathPatternVariant,
                                                   std::move(paths), addPathIndex, readSpec, computeActorId, retryPolicy,
-                                                  cfg, counters, taskCounters, fileSizeLimit, rowsLimitHint, memoryQuotaManager,
+                                                  cfg, counters, taskCounters, fileSizeLimit, sizeLimit, rowsLimitHint, memoryQuotaManager,
                                                   params.GetUseRuntimeListing(), fileQueueActor, fileQueueBatchSizeLimit, fileQueueBatchObjectCountLimit, fileQueueConsumersCountDelta);
 
         return {actor, actor};
diff --git a/ydb/library/yql/providers/s3/actors/yql_s3_read_actor.h b/ydb/library/yql/providers/s3/actors/yql_s3_read_actor.h
@@ -18,6 +18,7 @@ NActors::IActor* CreateS3FileQueueActor(
         NS3Details::TPathList paths,
         size_t prefetchSize,
         ui64 fileSizeLimit,
+        ui64 readLimit,
         bool useRuntimeListing,
         ui64 consumersCount,
         ui64 batchSizeLimit,
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_dq_integration.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_dq_integration.cpp
@@ -410,7 +410,7 @@ class TS3DqIntegration: public TDqIntegrationBase {
 
             auto fileQueueBatchObjectCountLimit = State_->Configuration->FileQueueBatchObjectCountLimit.Get().GetOrElse(1000);
             srcDesc.MutableSettings()->insert({"fileQueueBatchObjectCountLimit", ToString(fileQueueBatchObjectCountLimit)});
-            
+
             YQL_CLOG(DEBUG, ProviderS3) << " useRuntimeListing=" << useRuntimeListing;
 
             if (useRuntimeListing) {
@@ -422,8 +422,8 @@ class TS3DqIntegration: public TDqIntegrationBase {
                         packed.Data().Literal().Value(),
                         FromString<bool>(packed.IsText().Literal().Value()),
                         paths);
-                    paths.insert(paths.end(), 
-                        std::make_move_iterator(pathsChunk.begin()), 
+                    paths.insert(paths.end(),
+                        std::make_move_iterator(pathsChunk.begin()),
                         std::make_move_iterator(pathsChunk.end()));
                 }
 
@@ -434,11 +434,11 @@ class TS3DqIntegration: public TDqIntegrationBase {
                     builder.AddPath(f.Path, f.Size, f.IsDirectory);
                 });
                 builder.Save(&range);
-                
+
                 TVector<TString> serialized(1);
                 TStringOutput out(serialized.front());
                 range.Save(&out);
-                
+
                 paths.clear();
                 ReadPathsList(srcDesc, {}, serialized, paths);
 
@@ -485,12 +485,18 @@ class TS3DqIntegration: public TDqIntegrationBase {
 
                 YQL_CLOG(DEBUG, ProviderS3) << " hasDirectories=" << hasDirectories << ", consumersCount=" << consumersCount;
 
+                ui64 readLimit = std::numeric_limits<ui64>::max();
+                if (srcDesc.MutableSettings()->find("sizeLimit") != srcDesc.MutableSettings()->cend()) {
+                    readLimit = FromString<ui64>(srcDesc.MutableSettings()->at("sizeLimit"));
+                }
+
                 auto fileQueueActor = NActors::TActivationContext::ActorSystem()->Register(
                     NDq::CreateS3FileQueueActor(
                         0ul,
                         std::move(paths),
                         fileQueuePrefetchSize,
                         fileSizeLimit,
+                        readLimit,
                         useRuntimeListing,
                         consumersCount,
                         fileQueueBatchSizeLimit,
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
@@ -233,8 +233,15 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {
                                 fileSizeLimit = it->second;
                             }
                         }
+
+                        ui64 userSizeLimit = std::numeric_limits<ui64>::max();
                         if (formatName == "parquet") {
                             fileSizeLimit = State_->Configuration->BlockFileSizeLimit;
+                        } else if (formatName == "raw") {
+                            const auto sizeLimitParam = dqSource.Input().Cast<TS3SourceSettings>().SizeLimit().Maybe<TCoAtom>();
+                            if (sizeLimitParam.IsValid()) {
+                                userSizeLimit = FromString<ui64>(sizeLimitParam.Cast().StringValue());
+                            }
                         }
 
                         for (const TS3Path& batch : maybeS3SourceSettings.Cast().Paths()) {
@@ -245,13 +252,13 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {
                             UnpackPathsList(packed, isTextEncoded, paths);
 
                             for (auto& entry : paths) {
-                                if (entry.Size > fileSizeLimit) {
+                                if (std::min(entry.Size, userSizeLimit) > fileSizeLimit) {
                                     ctx.AddError(TIssue(ctx.GetPosition(batch.Pos()),
                                         TStringBuilder() << "Size of object " << entry.Path << " = " << entry.Size << " and exceeds limit = " << fileSizeLimit << " specified for format " << formatName));
                                     hasErr = true;
                                     return false;
                                 }
-                                totalSize += entry.Size;
+                                totalSize += std::min(entry.Size, userSizeLimit);
                                 ++count;
                             }
                         }
@@ -673,7 +680,7 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {
                     .RowsLimitHint(count.Literal())
                     .Build()
             .Build()
-        .Done();   
+        .Done();
     }
 
     TMaybeNode<TExprBase> PushDownLimit(TExprBase node, TExprContext& ctx) const {
diff --git a/ydb/tests/fq/s3/test_size_limit.py b/ydb/tests/fq/s3/test_size_limit.py

Original file line number	Diff line number	Diff line change
`@@ -233,8 +233,15 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {`
`233`	`233`	`fileSizeLimit = it->second;`
`234`	`234`	`}`
`235`	`235`	`}`
	`236`	`+`
	`237`	`+ ui64 userSizeLimit = std::numeric_limits<ui64>::max();`
`236`	`238`	`if (formatName == "parquet") {`
`237`	`239`	`fileSizeLimit = State_->Configuration->BlockFileSizeLimit;`
	`240`	`+ } else if (formatName == "raw") {`
	`241`	`+ const auto sizeLimitParam = dqSource.Input().Cast<TS3SourceSettings>().SizeLimit().Maybe<TCoAtom>();`
	`242`	`+ if (sizeLimitParam.IsValid()) {`
	`243`	`+ userSizeLimit = FromString<ui64>(sizeLimitParam.Cast().StringValue());`
	`244`	`+ }`
`238`	`245`	`}`
`239`	`246`
`240`	`247`	`for (const TS3Path& batch : maybeS3SourceSettings.Cast().Paths()) {`
`@@ -245,13 +252,13 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {`
`245`	`252`	`UnpackPathsList(packed, isTextEncoded, paths);`
`246`	`253`
`247`	`254`	`for (auto& entry : paths) {`
`248`		`- if (entry.Size > fileSizeLimit) {`
	`255`	`+ if (std::min(entry.Size, userSizeLimit) > fileSizeLimit) {`
`249`	`256`	`ctx.AddError(TIssue(ctx.GetPosition(batch.Pos()),`
`250`	`257`	`TStringBuilder() << "Size of object " << entry.Path << " = " << entry.Size << " and exceeds limit = " << fileSizeLimit << " specified for format " << formatName));`
`251`	`258`	`hasErr = true;`
`252`	`259`	`return false;`
`253`	`260`	`}`
`254`		`- totalSize += entry.Size;`
	`261`	`+ totalSize += std::min(entry.Size, userSizeLimit);`
`255`	`262`	`++count;`
`256`	`263`	`}`
`257`	`264`	`}`
`@@ -673,7 +680,7 @@ class TS3LogicalOptProposalTransformer : public TOptimizeTransformerBase {`
`673`	`680`	`.RowsLimitHint(count.Literal())`
`674`	`681`	`.Build()`
`675`	`682`	`.Build()`
`676`		`- .Done();`
	`683`	`+ .Done();`
`677`	`684`	`}`
`678`	`685`
`679`	`686`	`TMaybeNode<TExprBase> PushDownLimit(TExprBase node, TExprContext& ctx) const {`