Add flag for predicate extract for data queries

ref:8c53bd61f70e590a15d07e7fd32dd4b3ace4ee5f

7 files changed, 44 insertions, 5 deletions

diff --git a/ydb/core/kqp/opt/kqp_opt_kql.cpp b/ydb/core/kqp/opt/kqp_opt_kql.cpp
index aff06563c62..e36dbe7c7fc 100644
--- a/ydb/core/kqp/opt/kqp_opt_kql.cpp
+++ b/ydb/core/kqp/opt/kqp_opt_kql.cpp
@@ -30,6 +30,10 @@ bool UseReadTableRanges(const TKikimrTableDescription& tableData, const TIntrusi
     if (kqpCtx->IsScanQuery() && kqpCtx->Config->FeatureFlags.GetEnablePredicateExtractForScanQueries()) {
         return true;
     }
+    
+    if (kqpCtx->IsDataQuery() && kqpCtx->Config->FeatureFlags.GetEnablePredicateExtractForDataQueries()) {
+        return true;
+    }
 
     return false;
 }
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
index 0ce9229abe6..c0be2b4dc5b 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
@@ -207,7 +207,27 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
 
     auto rightReadMatch = MatchRead<TKqlReadTableBase>(join.RightInput());
     if (!rightReadMatch) {
-        return {};
+        if (auto readRangesMatch = MatchRead<TKqlReadTableRangesBase>(join.RightInput())) {
+            // for now only fullscans are supported
+            auto read = readRangesMatch->Read.Cast<TKqlReadTableRangesBase>();
+            if (TCoVoid::Match(read.Ranges().Raw())) {
+                rightReadMatch = readRangesMatch;
+                rightReadMatch->Read =
+                    Build<TKqlReadTable>(ctx, join.Pos())
+                    .Settings(read.Settings())
+                    .Table(read.Table())
+                    .Columns(read.Columns())
+                    .Range<TKqlKeyRange>()
+                        .From<TKqlKeyInc>().Build()
+                        .To<TKqlKeyInc>().Build()
+                        .Build()
+                    .Done();
+            } else {
+                return {};
+            }
+        } else {
+            return {};
+        }
     }
 
     if (rightReadMatch->FlatMap && !IsPassthroughFlatMap(rightReadMatch->FlatMap.Cast(), nullptr)) {
diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy_precompute.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy_precompute.cpp
index b009486741d..80998db221c 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy_precompute.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy_precompute.cpp
@@ -56,6 +56,12 @@ TExprBase KqpPropagatePrecomuteScalarRowset(TExprBase node, TExprContext& ctx, I
                     return false;
                 }
 
+                auto consumers = GetConsumers(maybePrecompute.Cast(), parentsMap);
+                if (consumers.size() > 1) {
+                    visitResult = false;
+                    return false;
+                }
+
                 precompute = maybePrecompute;
                 return false;
             }
diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.cpp b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
index 46b5de9186f..39dfd130e3e 100644
--- a/ydb/core/kqp/ut/common/kqp_ut_common.cpp
+++ b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
@@ -110,7 +110,7 @@ TKikimrRunner::TKikimrRunner(const TKikimrSettings& settings) {
     ServerSettings->SetKeepSnapshotTimeout(settings.KeepSnapshotTimeout);
     ServerSettings->SetFrFactory(&UdfFrFactory);
     ServerSettings->SetEnableNotNullColumns(true);
-    ServerSettings->SetEnableKqpScanQueryStreamLookup(true);
+    ServerSettings->SetEnableKqpScanQueryStreamLookup(false);
     if (settings.LogStream)
         ServerSettings->SetLogBackend(new TStreamLogBackend(settings.LogStream));
 
diff --git a/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp b/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
index 5c6066ea03f..f940b9cfc43 100644
--- a/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_perf_ut.cpp
@@ -565,7 +565,7 @@ Y_UNIT_TEST_SUITE(KqpPerf) {
         CompareYson(R"([[["Anna"];[3800u]]])", FormatResultSetYson(result.GetResultSet(0)));
 
         auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
-        UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(stats.query_phases().size(), UseNewEngine ? 2 : 1);
     }
 
     Y_UNIT_TEST_TWIN(MultiDeleteFromTable, UseNewEngine) {
diff --git a/ydb/core/kqp/ut/kqp_scan_ut.cpp b/ydb/core/kqp/ut/kqp_scan_ut.cpp
index 86a948ade19..bc430ba986a 100644
--- a/ydb/core/kqp/ut/kqp_scan_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_scan_ut.cpp
@@ -1941,7 +1941,11 @@ Y_UNIT_TEST_SUITE(KqpScan) {
     }
 
     Y_UNIT_TEST_TWIN(StreamLookup, UseSessionActor) {
-        auto kikimr = KikimrRunnerEnableSessionActor(UseSessionActor);
+        auto settings = TKikimrSettings()
+            .SetEnableKqpSessionActor(UseSessionActor)
+            .SetEnableKqpScanQueryStreamLookup(true);
+
+        TKikimrRunner kikimr(settings);
         auto db = kikimr.GetTableClient();
         CreateSampleTables(kikimr);
 
@@ -1970,7 +1974,11 @@ Y_UNIT_TEST_SUITE(KqpScan) {
     }
 
     Y_UNIT_TEST_TWIN(StreamLookupByPkPrefix, UseSessionActor) {
-        auto kikimr = KikimrRunnerEnableSessionActor(UseSessionActor);
+        auto settings = TKikimrSettings()
+            .SetEnableKqpSessionActor(UseSessionActor)
+            .SetEnableKqpScanQueryStreamLookup(true);
+
+        TKikimrRunner kikimr(settings);
         auto db = kikimr.GetTableClient();
         CreateSampleTables(kikimr);
 
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index cd79e3e913b..695ee5d85e0 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -701,6 +701,7 @@ message TFeatureFlags {
     optional bool EnableChangefeeds = 65 [default = true];
     optional bool EnableKqpScanQueryStreamLookup = 66 [default = false];
     optional bool EnableKqpScanQueryMultipleOlapShardsReads = 67 [default = false];
+    optional bool EnablePredicateExtractForDataQueries = 68 [default = false];
 }