Fix MaterializeSectionIfRequired for section with sampling, limits, and sys columns

commit_hash:14bce76d19bc8369d2d9f9ef863d966f70a59f91

3 files changed, 16 insertions, 2 deletions

diff --git a/yt/yql/providers/yt/provider/yql_yt_optimize.cpp b/yt/yql/providers/yt/provider/yql_yt_optimize.cpp
index 0137b343eb..cd0b5e2ade 100644
--- a/yt/yql/providers/yt/provider/yql_yt_optimize.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_optimize.cpp
@@ -73,7 +73,7 @@ TMaybeNode<TYtSection> MaterializeSectionIfRequired(TExprBase world, TYtSection
                 .Paths()
                     .Add(path)
                 .Build()
-                .Settings(NYql::RemoveSetting(section.Settings().Ref(), EYtSettingType::Sample, ctx))
+                .Settings(NYql::RemoveSettings(section.Settings().Ref(), EYtSettingType::Sample | EYtSettingType::SysColumns, ctx))
                 .Done();
     }
 
@@ -89,7 +89,7 @@ TMaybeNode<TYtSection> UpdateSectionWithRange(TExprBase world, TYtSection sectio
     TVector<TYtPath> skippedPaths;
     if (auto limiter = TTableLimiter(range)) {
         if (auto materialized = MaterializeSectionIfRequired(world, section, dataSink, outRowSpec, keepSortness,
-            {NYql::KeepOnlySettings(section.Settings().Ref(), EYtSettingType::Take | EYtSettingType::Skip | EYtSettingType::SysColumns, ctx)}, state, ctx))
+            {NYql::KeepOnlySettings(section.Settings().Ref(), EYtSettingType::Take | EYtSettingType::Skip, ctx)}, state, ctx))
         {
             if (!allowMaterialize || state->Types->EvaluationInProgress) {
                 // Keep section as is
diff --git a/yt/yql/tests/sql/suites/select/sample_limit_recordindex.cfg b/yt/yql/tests/sql/suites/select/sample_limit_recordindex.cfg
new file mode 100644
index 0000000000..6c06cba116
--- /dev/null
+++ b/yt/yql/tests/sql/suites/select/sample_limit_recordindex.cfg
@@ -0,0 +1,2 @@
+in Input input1100.txt
+
diff --git a/yt/yql/tests/sql/suites/select/sample_limit_recordindex.sql b/yt/yql/tests/sql/suites/select/sample_limit_recordindex.sql
new file mode 100644
index 0000000000..220fbfa060
--- /dev/null
+++ b/yt/yql/tests/sql/suites/select/sample_limit_recordindex.sql
@@ -0,0 +1,12 @@
+/* custom check: len(yt_res_yson[0][b'Write'][0][b'Data']) <= 5 */
+USE plato;
+
+SELECT
+    key,
+    subkey,
+    TableRecordIndex() AS index
+FROM
+    Input
+SAMPLE 1.0 / 5
+LIMIT 5
+;