Merge cdc885b into e2ed478

ydb-platform · Jul 2, 2024 · 27be527 · 27be527
2 parents e2ed478 + cdc885b
commit 27be527
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
@@ -2109,8 +2109,8 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
                 Finish(Ydb::StatusIds::INTERNAL_ERROR, "Result set row is empty");
                 return;
             }
-  
-            i64 rowSize = serializedRow->size();
+
+            i64 rowSize = serializedRow->size() + AdditionalRowSize;
             if (SizeLimit && ResultSet.rows_size() && ResultSetSize + rowSize > SizeLimit) {
                 CancelFetchQuery();
                 return;
@@ -2121,11 +2121,17 @@ class TGetScriptExecutionResultQueryActor : public TQueryBase {
                 return;
             }
 
-            ResultSetSize += rowSize;
             if (!ResultSet.add_rows()->ParseFromString(*serializedRow)) {
                 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 - rowSize;
+                rowSize += AdditionalRowSize;
+            }
+            ResultSetSize += rowSize;
         }
 
         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);
     }
 }