Fix sink empty batch (#10463)

ydb-platform · Oct 16, 2024 · 0025e00 · 0025e00
1 parent 6b4d585
commit 0025e00
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/ydb/core/kqp/runtime/kqp_write_table.cpp b/ydb/core/kqp/runtime/kqp_write_table.cpp
@@ -889,14 +889,17 @@ class TShardsInfo {
 
         void MakeNextBatches(i64 maxDataSize, ui64 maxCount) {
             YQL_ENSURE(BatchesInFlight == 0);
+            YQL_ENSURE(!IsEmpty());
             i64 dataSize = 0;
+            // For columnshard batch can be slightly larger than the limit.
             while (BatchesInFlight < maxCount
                     && BatchesInFlight < Batches.size()
-                    && dataSize + GetBatch(BatchesInFlight)->GetMemory() <= maxDataSize) {
+                    && (dataSize + GetBatch(BatchesInFlight)->GetMemory() <= maxDataSize || BatchesInFlight == 0)) {
                 dataSize += GetBatch(BatchesInFlight)->GetMemory();
                 ++BatchesInFlight;
             }
-            YQL_ENSURE(BatchesInFlight == Batches.size() || GetBatch(BatchesInFlight)->GetMemory() <= maxDataSize); 
+            YQL_ENSURE(BatchesInFlight != 0);
+            YQL_ENSURE(BatchesInFlight == maxCount || BatchesInFlight == Batches.size() || dataSize + GetBatch(BatchesInFlight)->GetMemory() >= maxDataSize);
         }
 
         const IPayloadSerializer::IBatchPtr& GetBatch(size_t index) const {
@@ -1200,7 +1203,9 @@ class TShardedWriteController : public IShardedWriteController {
         if (force) {
             for (auto& [shardId, batches] : Serializer->FlushBatchesForce()) {
                 for (auto& batch : batches) {
-                    ShardsInfo.GetShard(shardId).PushBatch(std::move(batch));
+                    if (batch && !batch->IsEmpty()) {
+                        ShardsInfo.GetShard(shardId).PushBatch(std::move(batch));
+                    }
                 }
             }
         } else {