Drop aux columns that are not expected after row spec rebuild in AssumeConstraints

commit_hash:2326ce8e0746c8c27bc555fc4cb552116a6da32d

Author: loochek

yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_sort.cpp

@@ -667,11 +667,60 @@ TMaybeNode<TExprBase> TYtPhysicalOptProposalTransformer::AssumeConstraints(TExpr
         const size_t index = FromString(input.Cast<TYtOutput>().OutIndex().Value());
         TYtOutTableInfo outTable(op.Output().Item(index));
         if (builder) {
+            YQL_ENSURE(!builder->NeedMap());
             builder->FillRowSpecSort(*outTable.RowSpec);
         }
         outTable.RowSpec->SetConstraints(assume.Ref().GetConstraintSet());
         outTable.SetUnique(assume.Ref().GetConstraint<TDistinctConstraintNode>(), assume.Pos(), ctx);
 
+        if (op.Maybe<TYtMap>() || op.Maybe<TYtReduce>()) {
+            TExprNode::TPtr lambda;
+            size_t childToReplace = 0;
+            if (auto map = op.Maybe<TYtMap>()) {
+                lambda = map.Cast().Mapper().Ptr();
+                childToReplace = TYtMap::idx_Mapper;
+            } else if (auto reduce = op.Maybe<TYtReduce>()) {
+                lambda = reduce.Cast().Reducer().Ptr();
+                childToReplace = TYtReduce::idx_Reducer;
+            } else {
+                YQL_ENSURE(false, "unexpected operation");
+            }
+
+            auto actualLambdaOutputType = TCoLambda(lambda).Body().Ref().GetTypeAnn();
+            auto expectedLambdaOutputType = outTable.RowSpec->GetExtendedType(ctx);
+            if (!IsSameAnnotation(*actualLambdaOutputType, *expectedLambdaOutputType)) {
+                // Drop aux columns that are not expected after row spec rebuild
+
+                auto excludeLambda = Build<TCoLambda>(ctx, lambda->Pos())
+                    .Args({"stream"})
+                    .Body<TCoOrderedMap>()
+                        .Input("stream")
+                        .Lambda<TCoLambda>()
+                            .Args({"struct"})
+                            .Body<TCoCastStruct>()
+                                .Struct("struct")
+                                .Type(ExpandType(lambda->Pos(), *expectedLambdaOutputType, ctx))
+                            .Build()
+                        .Build()
+                    .Build()
+                    .Done();
+
+                lambda = Build<TCoLambda>(ctx, lambda->Pos())
+                    .Args({"stream"})
+                    .Body<TExprApplier>()
+                        .Apply(excludeLambda)
+                        .With<TExprApplier>(0)
+                            .Apply(TCoLambda(lambda))
+                            .With(0, "stream")
+                        .Build()
+                    .Build()
+                    .Done()
+                    .Ptr();
+
+                newOp = ctx.ChangeChild(*newOp, childToReplace, std::move(lambda));
+            }
+        }
+
         TVector<TYtOutTable> outputs;
         for (size_t i = 0; i < op.Output().Size(); ++i) {
             if (index == i) {

yt/yql/tests/sql/suites/order_by/input_sorted_desc.txt

@@ -0,0 +1,3 @@
+{"_yql_column_0"="\xE0\xCC\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFE";"key2"=3;"key1"=3;};
+{"_yql_column_0"="\xE0\xCD\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFE";"key2"=2;"key1"=2;};
+{"_yql_column_0"="\xE0\xCE\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFE";"key2"=1;"key1"=1;};

yt/yql/tests/sql/suites/order_by/input_sorted_desc.txt.attr

@@ -0,0 +1,61 @@
+{
+    "_yql_row_spec" = {
+        "Constraints" = {
+            "Sorted" = [
+                [
+                    [
+                        "key1";
+                    ];
+                    %false;
+                ];
+                [
+                    [
+                        "key2";
+                    ];
+                    %true;
+                ];
+            ];
+        };
+        "SortDirections" = [
+            0;
+            1;
+        ];
+        "SortMembers" = [
+            "key1";
+            "key2";
+        ];
+        "SortedBy" = [
+            "_yql_column_0";
+            "key2";
+        ];
+        "SortedByTypes" = [
+            [
+                "DataType";
+                "String";
+            ];
+            [
+                "DataType";
+                "Int32";
+            ];
+        ];
+        "Type" = [
+            "StructType";
+            [
+                [
+                    "key1";
+                    [
+                        "DataType";
+                        "Int32";
+                    ];
+                ];
+                [
+                    "key2";
+                    [
+                        "DataType";
+                        "Int32";
+                    ];
+                ];
+            ];
+        ];
+    };
+}

yt/yql/tests/sql/suites/order_by/yql-19598.cfg

@@ -0,0 +1 @@
+in Input input_sorted_desc.txt

yt/yql/tests/sql/suites/order_by/yql-19598.sql

@@ -0,0 +1,4 @@
+/* postgres can not */
+USE plato;
+
+SELECT * FROM Input WHERE key2 = 2 ASSUME ORDER BY key2;