fix-write

Author: nikvas0

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,18 +1203,19 @@ 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 {
             for (const ui64 shardId : Serializer->GetShardIds()) {
                 auto& shard = ShardsInfo.GetShard(shardId);
                 while (true) {
                     auto batch = Serializer->FlushBatch(shardId);
-                    if (!batch || batch->IsEmpty()) {
-                        break;
+                    if (batch && !batch->IsEmpty()) {
+                        shard.PushBatch(std::move(batch));
                     }
-                    shard.PushBatch(std::move(batch));
                 }
             }
         }