Merge 8932c0e into c7a7fef

ydb-platform · Jul 2, 2024 · a490e71 · a490e71
2 parents c7a7fef + 8932c0e
commit a490e71
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/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;
 };

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
@@ -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);
     }
 }