1 files changed, 41 insertions, 3 deletions

diff --git a/yql/essentials/core/yql_expr_constraint.cpp b/yql/essentials/core/yql_expr_constraint.cpp
index 864f46e248d..eb8a9694f63 100644
--- a/yql/essentials/core/yql_expr_constraint.cpp
+++ b/yql/essentials/core/yql_expr_constraint.cpp
@@ -207,6 +207,7 @@ public:
         Functions_["DynamicVariant"] = &TCallableConstraintTransformer::DynamicVariantWrap;
         Functions_["Guess"] = &TCallableConstraintTransformer::GuessWrap;
         Functions_["Mux"] = &TCallableConstraintTransformer::MuxWrap;
+        Functions_["Demux"] = &TCallableConstraintTransformer::FromFirst<TStreamingConstraintNode>;
         Functions_["Nth"] = &TCallableConstraintTransformer::NthWrap;
         Functions_["EquiJoin"] = &TCallableConstraintTransformer::EquiJoinWrap;
         Functions_["JoinDict"] = &TCallableConstraintTransformer::JoinDictWrap;
@@ -1844,6 +1845,8 @@ private:
     TStatus SwitchWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
         TStatus status = TStatus::Ok;
         TDynBitMap outFromChildren; // children, from which take a multi constraint for output
+        const auto* inputStreaming = input->Head().GetConstraint<TStreamingConstraintNode>();
+
         if (const auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
             for (size_t i = 2; i < input->ChildrenSize(); ++i) {
                 TMultiConstraintNode::TMapType items;
@@ -1861,18 +1864,33 @@ private:
                     if (input->Child(i)->ChildrenSize() > 1) {
                         argConstraints.push_back(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items)));
                         argConstraints.push_back(ctx.MakeConstraint<TVarIndexConstraintNode>(input->Child(i)->ChildrenSize()));
+
+                        if (inputStreaming) {
+                            argConstraints.push_back(inputStreaming);
+                        }
                     } else {
                         argConstraints = items.front().second.GetAllConstraints();
+
+                        if (inputStreaming && !items.front().second.GetConstraint<TStreamingConstraintNode>()) {
+                            argConstraints.push_back(inputStreaming);
+                        }
                     }
+                } else if (inputStreaming) {
+                    argConstraints.push_back(inputStreaming);
                 }
 
                 status = status.Combine(UpdateLambdaConstraints(input->ChildRef(++i), ctx, {argConstraints}));
             }
         } else {
-            const bool inVar = GetSeqItemType(*input->Head().GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant;
-            const TSmallVec<TConstraintNode::TListType> argConstraints(1U, inVar ? TConstraintNode::TListType() : input->Head().GetAllConstraints());
+            TConstraintNode::TListType argConstraintsList;
+            if (GetSeqItemType(*input->Head().GetTypeAnn()).GetKind() != ETypeAnnotationKind::Variant) {
+                argConstraintsList = input->Head().GetAllConstraints();
+            } else if (inputStreaming) {
+                argConstraintsList.push_back(inputStreaming);
+            }
+
             for (size_t i = 3; i < input->ChildrenSize(); i += 2) {
-                status = status.Combine(UpdateLambdaConstraints(input->ChildRef(i), ctx, argConstraints));
+                status = status.Combine(UpdateLambdaConstraints(input->ChildRef(i), ctx, {argConstraintsList}));
             }
             outFromChildren.Set(0, input->ChildrenSize());
         }
@@ -1883,6 +1901,7 @@ private:
 
         const auto inputVarIndex = input->Head().GetConstraint<TVarIndexConstraintNode>();
         const bool emptyInput = input->Head().GetConstraint<TEmptyConstraintNode>();
+        bool hasStreamingOutput = false;
         if (GetSeqItemType(*input->GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant) {
             ui32 outIndexOffset = 0;
             TMultiConstraintNode::TMapType multiItems;
@@ -1948,6 +1967,9 @@ private:
             }
 
             if (!multiItems.empty()) {
+                hasStreamingOutput = std::any_of(multiItems.begin(), multiItems.end(), [](const auto& item) {
+                    return !!item.second.template GetConstraint<TStreamingConstraintNode>();
+                });
                 input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems)));
             }
             if (emptyOut) {
@@ -1957,6 +1979,11 @@ private:
             YQL_ENSURE(input->ChildrenSize() == 4);
             input->CopyConstraints(*input->Child(3));
         }
+
+        if (hasStreamingOutput || inputStreaming) {
+            input->AddConstraint(ctx.MakeConstraint<TStreamingConstraintNode>());
+        }
+
         return FromFirst<TEmptyConstraintNode>(input, output, ctx);
     }
 
@@ -2234,19 +2261,26 @@ private:
                 TMultiConstraintNode::TMapType items;
                 ui32 index = 0;
                 ui32 emptyCount = 0;
+                bool hasStreaming = false;
                 for (auto& child: input->Head().Children()) {
                     if (!child->GetConstraint<TEmptyConstraintNode>()) {
                         items.push_back(std::make_pair(index, child->GetConstraintSet()));
+                        hasStreaming = hasStreaming || !!child->GetConstraint<TStreamingConstraintNode>();
                     } else {
                         ++emptyCount;
                     }
                     ++index;
                 }
+
                 if (!items.empty()) {
                     input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items)));
                 } else if (index == emptyCount) {
                     input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
                 }
+
+                if (hasStreaming) {
+                    input->AddConstraint(ctx.MakeConstraint<TStreamingConstraintNode>());
+                }
             }
         }
         return TStatus::Ok;
@@ -2289,6 +2323,10 @@ private:
                     input->SetConstraints(*c);
                 }
             }
+
+            if (const auto* c = input->Head().GetConstraint<TStreamingConstraintNode>()) {
+                input->AddConstraint(c);
+            }
         }
         return TStatus::Ok;
     }