Merge 5d2dd8b into 2d3caea

Author: lll-phill-lll

ydb/library/yql/minikql/comp_nodes/mkql_wide_combine.cpp

@@ -132,6 +132,24 @@ struct TCombinerNodes {
         }
     }
 
+    void ExtractValues(TComputationContext& ctx, NUdf::TUnboxedValue** from, NUdf::TUnboxedValue* to) const {
+        for (ui32 i = 0U; i < ItemNodes.size(); ++i) {
+            if (from[i]) {
+                to[i] = std::move(*(from[i]));
+            }
+        }
+    }
+
+    void ExtractValues(TComputationContext& ctx, NUdf::TUnboxedValue* from, NUdf::TUnboxedValue** to) const {
+        for (size_t i = 0, j = 0; i != ItemNodes.size(); ++i) {
+            if (IsInputItemNodeUsed(i)) {
+                *to[i] = std::move(from[j++]);
+            } else {
+                to[i] = nullptr;
+            }
+        }
+    }
+
     void ProcessItem(TComputationContext& ctx, NUdf::TUnboxedValue* keys, NUdf::TUnboxedValue* state) const {
         if (keys) {
             std::fill_n(keys, KeyResultNodes.size(), NUdf::TUnboxedValuePod());
@@ -346,16 +364,17 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     enum class ETasteResult: i8 {
         Init = -1,
         Update,
+        DataRequired,
+        RawDataExtractionRequired,
         Skip
     };
     TSpillingSupportState(
-        TMemoryUsageInfo* memInfo, size_t wideFieldsIndex,
+        TMemoryUsageInfo* memInfo,
         const TMultiType* usedInputItemType, const TMultiType* keyAndStateType, ui32 keyWidth, size_t itemNodesSize,
         const THashFunc& hash, const TEqualsFunc& equal, bool allowSpilling, TComputationContext& ctx
     )
         : TBase(memInfo)
         , InMemoryProcessingState(memInfo, keyWidth, keyAndStateType->GetElementsCount() - keyWidth, hash, equal)
-        , WideFieldsIndex(wideFieldsIndex)
         , UsedInputItemType(usedInputItemType)
         , KeyAndStateType(keyAndStateType)
         , KeyWidth(keyWidth)
@@ -380,7 +399,7 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     bool IsProcessingRequired() const {
         if (InputStatus != EFetchResult::Finish) return true;
 
-        return HasDataForProcessing;
+        return HasRawDataToExtract || HasDataForProcessing;
     }
 
     bool UpdateAndWait() {
@@ -424,10 +443,19 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
             return isNew ? ETasteResult::Init : ETasteResult::Update;
         }
         if (GetMode() == EOperatingMode::ProcessSpilled) {
+            if (HasRawDataToExtract) {
+                // Tongue not used here.
+                Throat = BufferForUsedInputItems.data();
+                HasRawDataToExtract = false;
+                HasDataForProcessing = true;
+                return ETasteResult::RawDataExtractionRequired;
+            }
+            HasDataForProcessing = false;
             // while restoration we process buckets one by one starting from the first in a queue
             bool isNew = SpilledBuckets.front().InMemoryProcessingState->TasteIt();
             Throat = SpilledBuckets.front().InMemoryProcessingState->Throat;
             Tongue = SpilledBuckets.front().InMemoryProcessingState->Tongue;
+            BufferForUsedInputItems.resize(0);
             return isNew ? ETasteResult::Init : ETasteResult::Update;
         }
 
@@ -445,9 +473,13 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
 
         // Corresponding bucket is spilled, we don't need a key anymore, full input will be spilled
         BufferForKeyAndState.resize(0);
-        TryToSpillRawData(bucket, bucketId);
+        // Prepare space for raw data 
+        MKQL_ENSURE(BufferForUsedInputItems.size() == 0, "Internal logic error");
+        BufferForUsedInputItems.resize(ItemNodesSize);
+        BufferForUsedInputItemsBucketId = bucketId;
+        Throat = BufferForUsedInputItems.data();
 
-        return ETasteResult::Skip;
+        return ETasteResult::DataRequired;
     }
 
     NUdf::TUnboxedValuePod* Extract() {
@@ -472,25 +504,6 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
         BufferForKeyAndState.resize(0);
     }
 
-    // Copies data from WideFields to local and tries to spill it using suitable bucket.
-    // if the bucket is already busy, then the buffer will wait for the next iteration.
-    void TryToSpillRawData(TSpilledBucket& bucket, size_t bucketId) {
-        auto **fields = Ctx.WideFields.data() + WideFieldsIndex;
-        MKQL_ENSURE(BufferForUsedInputItems.empty(), "Internal logic error");
-
-        for (size_t i = 0; i < ItemNodesSize; ++i) {
-            if (fields[i]) {
-                BufferForUsedInputItems.push_back(*fields[i]);
-            }
-        }
-        if (bucket.AsyncWriteOperation.has_value()) {
-            BufferForUsedInputItemsBucketId = bucketId;
-            return;
-        }
-        bucket.AsyncWriteOperation = bucket.SpilledData->WriteWideItem(BufferForUsedInputItems);
-        BufferForUsedInputItems.resize(0);
-    }
-
     bool FlushSpillingBuffersAndWait() {
         UpdateSpillingBuckets();
 
@@ -620,8 +633,14 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
             }
             AsyncReadOperation = std::nullopt;
         }
+
         auto& bucket = SpilledBuckets.front();
         if (bucket.BucketState == TSpilledBucket::EBucketState::InMemory) return false;
+        if (HasDataForProcessing) {
+            Tongue = bucket.InMemoryProcessingState->Tongue;
+            Throat = bucket.InMemoryProcessingState->Throat;
+            return false;
+        }
         //recover spilled state
         while(!bucket.SpilledState->Empty()) {
             RecoverState = true;
@@ -651,17 +670,11 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
             if (AsyncReadOperation) {
                 return true;
             }
-            auto **fields = Ctx.WideFields.data() + WideFieldsIndex;
-            for (size_t i = 0, j = 0; i < ItemNodesSize; ++i) {
-                if (fields[i]) {
-                    fields[i] = &(BufferForUsedInputItems[j++]);
-                }
-            }
 
             Tongue = bucket.InMemoryProcessingState->Tongue;
             Throat = bucket.InMemoryProcessingState->Throat;
 
-            HasDataForProcessing = true;
+            HasRawDataToExtract = true;
             return false;
         }
         bucket.BucketState = TSpilledBucket::EBucketState::InMemory;
@@ -725,8 +738,9 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
 
     bool HasDataForProcessing = false;
 
+    bool HasRawDataToExtract = false;
+
     TState InMemoryProcessingState;
-    const size_t WideFieldsIndex;
     const TMultiType* const UsedInputItemType;
     const TMultiType* const KeyAndStateType;
     const size_t KeyWidth;
@@ -1237,6 +1251,7 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideLastCombinerWra
         , AllowSpilling(allowSpilling)
     {}
 
+    // MARK: DoCAlculate
     EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
         if (!state.HasValue()) {
             MakeState(ctx, state);
@@ -1246,14 +1261,12 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideLastCombinerWra
             auto **fields = ctx.WideFields.data() + WideFieldsIndex;
 
             while (true) {
-                for (auto i = 0U; i < Nodes.ItemNodes.size(); ++i)
-                    fields[i] = Nodes.GetUsedInputItemNodePtrOrNull(ctx, i);
-
                 if (ptr->UpdateAndWait()) {
                     return EFetchResult::Yield;
                 }
-
                 if (ptr->InputStatus != EFetchResult::Finish) {
+                    for (auto i = 0U; i < Nodes.ItemNodes.size(); ++i)
+                        fields[i] = Nodes.GetUsedInputItemNodePtrOrNull(ctx, i);
                     switch (ptr->InputStatus = Flow->FetchValues(ctx, fields)) {
                         case EFetchResult::One:
                             break;
@@ -1274,6 +1287,13 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideLastCombinerWra
                         case TSpillingSupportState::ETasteResult::Update:
                             Nodes.ProcessItem(ctx, static_cast<NUdf::TUnboxedValue*>(ptr->Tongue), static_cast<NUdf::TUnboxedValue*>(ptr->Throat));
                             break;
+                        case TSpillingSupportState::ETasteResult::DataRequired:
+                            Nodes.ExtractValues(ctx, fields, static_cast<NUdf::TUnboxedValue*>(ptr->Throat));
+                            break;
+                        case TSpillingSupportState::ETasteResult::RawDataExtractionRequired:
+                            Nodes.ExtractValues(ctx, static_cast<NUdf::TUnboxedValue*>(ptr->Throat), fields);
+                            break;
+
                         case TSpillingSupportState::ETasteResult::Skip:
                             break;
                     }
@@ -1553,7 +1573,7 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideLastCombinerWra
 #endif
 private:
     void MakeState(TComputationContext& ctx, NUdf::TUnboxedValue& state) const {
-        state = ctx.HolderFactory.Create<TSpillingSupportState>(WideFieldsIndex,
+        state = ctx.HolderFactory.Create<TSpillingSupportState>(
             UsedInputItemType, KeyAndStateType,
             Nodes.KeyNodes.size(),
             Nodes.ItemNodes.size(),