YYQ-3398 fixed result set size counting (#6157)

Author: GrigoriyPA

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:

ydb/core/kqp/proxy_service/kqp_script_executions.cpp

@@ -2109,9 +2109,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;
             }
@@ -2126,6 +2126,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) {
@@ -2161,6 +2167,7 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
     const TInstant Deadline;
 
     i64 ResultSetSize = 0;
+    i64 AdditionalRowSize = 0;
     Ydb::ResultSet ResultSet;
     bool HasMoreResults = false;
 };

ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp

@@ -1786,6 +1786,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);
     }
 }