Fix off-by-one error in variant index calculation

commit_hash:ecdf3e7f584281b4f004f6b3b31d45127dabb561

5 files changed, 61 insertions, 2 deletions

diff --git a/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_fuse.cpp b/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_fuse.cpp
index 1ddb248e4ec..71ed0e311f4 100644
--- a/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_fuse.cpp
+++ b/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_fuse.cpp
@@ -518,11 +518,16 @@ TMaybeNode<TExprBase> TYtPhysicalOptProposalTransformer::FuseReduceWithTrivialMa
                 .Callable("Visit")
                     .Arg(0, "item")
                     .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                        const bool inserted = newInput.size() > origReduceInputs;
                         for (size_t i = 0; i < newInput.size(); ++i) {
                             TString paramName = TStringBuilder() << "alt" << i;
                             TString remappedName = TStringBuilder() << "remapped" << i;
                             if (i != fusedMap->InputIndex) {
+                                size_t origInputIndex = i;
+                                if (i > fusedMap->InputIndex) {
+                                    size_t delta = newInput.size() - origReduceInputs;
+                                    YQL_ENSURE(i >= delta);
+                                    origInputIndex = i - delta;
+                                }
                                 parent
                                     .Atom(2 * i + 1, i)
                                     .Lambda(2 * i + 2)
@@ -533,7 +538,7 @@ TMaybeNode<TExprBase> TYtPhysicalOptProposalTransformer::FuseReduceWithTrivialMa
                                                     parent
                                                         .Callable(0, "Variant")
                                                             .Arg(0, paramName)
-                                                            .Atom(1, (i > fusedMap->InputIndex && inserted) ? i - 1 : i)
+                                                            .Atom(1, origInputIndex)
                                                             .Add(2, origVariantType)
                                                         .Seal();
                                                 } else {
diff --git a/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.cfg b/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.cfg
new file mode 100644
index 00000000000..d9507518791
--- /dev/null
+++ b/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.cfg
@@ -0,0 +1,4 @@
+in Input1 sorted1.txt
+in Input2 sorted2.txt
+in Input3 sorted3.txt
+providers yt
diff --git a/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.sql b/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.sql
new file mode 100644
index 00000000000..edc1ff7c195
--- /dev/null
+++ b/yt/yql/tests/sql/suites/produce/reduce_with_trivial_remaps2.sql
@@ -0,0 +1,29 @@
+use plato;
+pragma EmitUnionMerge;
+
+
+use plato;
+pragma warning("disable", "4510");
+
+$process_stream = ($stream) -> {
+   $init = ($item) -> (AsStruct(1u as cnt, $item as row));
+   $switch = ($_item, $_state) -> (false);
+   $update = ($item, $state) -> (AsStruct($state.cnt + 1u as cnt,
+       if(($item.value > $state.row.value) ?? false, $item, $state.row) as row));
+   $state = YQL::Collect(YQL::Condense1($stream, $init, $switch, $update));
+   return $state;
+};
+
+$udf = ($_key, $variant_stream) -> {
+   $itentity = ($x)->($x);
+   return $process_stream(YQL::OrderedMap($variant_stream, ($item)->(YQL::Visit($item, AsAtom('0'), $itentity, AsAtom('1'), $itentity))));
+};
+
+$src1 = select * from Input1 assume order by key, subkey;
+$src2 = select key, subkey, cast(value as String) as value from Input2 assume order by key, subkey;
+$src3 = select * from Input3 assume order by key, subkey;
+
+$src = select * from $src1 union all select * from $src2 union all select * from $src3;
+
+REDUCE $src, $src1 presort subkey ON key USING $udf(TableRow());
+
diff --git a/yt/yql/tests/sql/suites/produce/sorted3.txt b/yt/yql/tests/sql/suites/produce/sorted3.txt
new file mode 100644
index 00000000000..b214aab0d92
--- /dev/null
+++ b/yt/yql/tests/sql/suites/produce/sorted3.txt
@@ -0,0 +1,10 @@
+{"key"="023";"subkey"="3";"value"="aaa"};
+{"key"="037";"subkey"="5";"value"="ddd"};
+{"key"="075";"subkey"="1";"value"="abc"};
+{"key"="150";"subkey"="1";"value"="aaa"};
+{"key"="150";"subkey"="3";"value"="iii"};
+{"key"="150";"subkey"="8";"value"="zzz"};
+{"key"="200";"subkey"="7";"value"="qqq"};
+{"key"="527";"subkey"="4";"value"="bbb"};
+{"key"="761";"subkey"="6";"value"="ccc"};
+{"key"="911";"subkey"="2";"value"="kkk"};
diff --git a/yt/yql/tests/sql/suites/produce/sorted3.txt.attr b/yt/yql/tests/sql/suites/produce/sorted3.txt.attr
new file mode 100644
index 00000000000..391c1a05f6b
--- /dev/null
+++ b/yt/yql/tests/sql/suites/produce/sorted3.txt.attr
@@ -0,0 +1,11 @@
+{"_yql_row_spec"={
+    "Type"=["StructType";[
+        ["key";["DataType";"String"]];
+        ["subkey";["DataType";"String"]];
+        ["value";["DataType";"String"]]
+    ]];
+    "SortDirections"=[1;1;1;];
+    "SortedBy"=["key";"subkey";"value";];
+    "SortedByTypes"=[["DataType";"String";];["DataType";"String";];["DataType";"String";];];
+    "SortMembers"=["key";"subkey";"value";];
+}}