Merge f1d16fa into ee02567

Author: lll-phill-lll

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

@@ -132,6 +132,18 @@ struct TCombinerNodes {
         }
     }
 
+    void ExtractValues(TComputationContext& ctx, NUdf::TUnboxedValue** values, NUdf::TUnboxedValue* keys) const {
+        for (size_t i = 0U; i < ItemNodes.size(); ++i) {
+            keys[i] = std::move(*(values[i]));
+        }
+    }
+
+    void ExtractValues(TComputationContext& ctx, NUdf::TUnboxedValue* keys, NUdf::TUnboxedValue** values) const {
+        for (size_t i = 0U, j = 0U; i < ItemNodes.size(); ++i) {
+            keys[i] = std::move(*(values[j++]));
+        }
+    }
+
     void ProcessItem(TComputationContext& ctx, NUdf::TUnboxedValue* keys, NUdf::TUnboxedValue* state) const {
         if (keys) {
             std::fill_n(keys, KeyResultNodes.size(), NUdf::TUnboxedValuePod());
@@ -346,6 +358,8 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     enum class ETasteResult: i8 {
         Init = -1,
         Update,
+        DataRequired,
+        DataExtractionRequired,
         Skip
     };
     TSpillingSupportState(
@@ -445,9 +459,11 @@ 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 
+        BufferForUsedInputItems.resize(ItemNodesSize);
+        Throat = BufferForUsedInputItems.data();
 
-        return ETasteResult::Skip;
+        return ETasteResult::DataRequired;
     }
 
     NUdf::TUnboxedValuePod* Extract() {
@@ -620,8 +636,15 @@ 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;
+            HasDataForProcessing = false;
+            return false;
+        }
         //recover spilled state
         while(!bucket.SpilledState->Empty()) {
             RecoverState = true;
@@ -651,15 +674,14 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
             if (AsyncReadOperation) {
                 return true;
             }
-            auto **fields = Ctx.WideFields.data() + WideFieldsIndex;
+            /* 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;
+            Throat = BufferForKeyAndState.data();
 
             HasDataForProcessing = true;
             return false;
@@ -1237,6 +1259,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);
@@ -1274,6 +1297,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::DataExtractionRequired:
+                            Nodes.ExtractValues(ctx, static_cast<NUdf::TUnboxedValue*>(ptr->Throat), fields);
+                            break;
+
                         case TSpillingSupportState::ETasteResult::Skip:
                             break;
                     }