Merge 4c4c5c2 into 736fd80

ydb-platform · Jul 2, 2024 · fda8e90 · fda8e90
2 parents 736fd80 + 4c4c5c2
commit fda8e90
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/ydb/core/grpc_services/query/rpc_fetch_script_results.cpp b/ydb/core/grpc_services/query/rpc_fetch_script_results.cpp
@@ -25,7 +25,7 @@ using TEvFetchScriptResultsRequest = TGrpcRequestNoOperationCall<Ydb::Query::Fet
     Ydb::Query::FetchScriptResultsResponse>;
 
 constexpr i64 MAX_ROWS_LIMIT = 1000;
-constexpr i64 MAX_SIZE_LIMIT = 60_MB;
+constexpr i64 MAX_SIZE_LIMIT = 50_MB;
 
 class TFetchScriptResultsRPC : public TRpcRequestActor<TFetchScriptResultsRPC, TEvFetchScriptResultsRequest, false> {
 public:

diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
@@ -2157,9 +2157,9 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
                 Finish(Ydb::StatusIds::INTERNAL_ERROR, "Result set row is empty");
                 return;
             }
-  
+
             i64 rowSize = serializedRow->size();
-            if (SizeLimit && ResultSet.rows_size() && ResultSetSize + rowSize > SizeLimit) {
+            if (SizeLimit && ResultSet.rows_size() && ResultSetSize + rowSize + AdditionalRowSize > SizeLimit) {
                 CancelFetchQuery();
                 return;
             }
@@ -2174,6 +2174,12 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
                 Finish(Ydb::StatusIds::INTERNAL_ERROR, "Result set row is corrupted");
                 return;
             }
+
+            // Initialize AdditionalRowSize
+            if (ResultSet.rows_size() == 1) {
+                AdditionalRowSize = static_cast<i64>(ResultSet.ByteSizeLong()) - ResultSetSize;
+            }
+            ResultSetSize += AdditionalRowSize;
         }
 
         if (TInstant::Now() + TDuration::Seconds(5) + GetAverageTime() >= Deadline) {
@@ -2209,6 +2215,7 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
     const TInstant Deadline;
 
     i64 ResultSetSize = 0;
+    i64 AdditionalRowSize = 0;
     Ydb::ResultSet ResultSet;
     bool HasMoreResults = false;
 };

diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp
@@ -1774,6 +1774,10 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
     }
 
     Y_UNIT_TEST(ExecuteScriptWithLargeFile) {
+        ExecuteSelectQuery("test_bucket_execute_script_with_large_file", 65_MB, 500000);
+    }
+
+    Y_UNIT_TEST(ExecuteScriptWithThinFile) {
         ExecuteSelectQuery("test_bucket_execute_script_with_large_file", 5_MB, 500000);
     }
 }