use old optimizer for equiranges

4 files changed, 81 insertions, 3 deletions

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
index 8c4405e689..f254535a3c 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -25,7 +25,6 @@ public:
         , Config(config)
     {
 #define HNDL(name) "KqpLogical-"#name, Hndl(&TKqpLogicalOptTransformer::name)
-        AddHandler(0, &TCoFlatMap::Match, HNDL(PushExtractedPredicateToReadTable));
         AddHandler(0, &TCoFlatMap::Match, HNDL(PushPredicateToReadTable));
         AddHandler(0, &TCoAggregate::Match, HNDL(RewriteAggregate));
         AddHandler(0, &TCoTake::Match, HNDL(RewriteTakeSortToTopSort));
@@ -46,6 +45,7 @@ public:
         AddHandler(0, &TKqpReadOlapTableRangesBase::Match, HNDL(ApplyExtractMembersToReadOlapTable<false>));
         AddHandler(0, &TKqlLookupTableBase::Match, HNDL(ApplyExtractMembersToLookupTable<false>));
 
+        AddHandler(1, &TCoFlatMap::Match, HNDL(PushExtractedPredicateToReadTable));
         AddHandler(1, &TKqlReadTableIndex::Match, HNDL(RewriteIndexRead));
         AddHandler(1, &TKqlLookupIndex::Match, HNDL(RewriteLookupIndex));
         AddHandler(1, &TKqlStreamLookupIndex::Match, HNDL(RewriteStreamLookupIndex));
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
index beb2a12705..88e9a831fa 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges.cpp
@@ -116,6 +116,19 @@ TKqlKeyRange BuildKeyRangeExpr(const TKeyRange& keyRange, const TKikimrTableDesc
         .Done();
 }
 
+bool IsPointPrefix(const TKeyRange& range) {
+    size_t prefixLen = 0;
+    for (size_t i = 0; i < range.GetColumnRangesCount(); ++i) {
+        if (range.GetColumnRange(i).IsPoint() && i == prefixLen) {
+            prefixLen += 1;
+        }
+        if (range.GetColumnRange(i).IsDefined() && i >= prefixLen) {
+            return false;
+        }
+    }
+    return prefixLen > 0;
+}
+
 TExprBase KqpPushPredicateToReadTable(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx) {
     if (!node.Maybe<TCoFlatMap>()) {
         return node;
@@ -126,7 +139,35 @@ TExprBase KqpPushPredicateToReadTable(TExprBase node, TExprContext& ctx, const T
         return node;
     }
 
+    bool onlyPointRanges = false;
     auto readMatch = MatchRead<TKqlReadTableBase>(flatmap.Input());
+
+    //TODO: remove this branch KIKIMR-15255, KIKIMR-15321
+    if (!readMatch && kqpCtx.IsDataQuery()) {
+        if (auto readRangesMatch = MatchRead<TKqlReadTableRangesBase>(flatmap.Input())) {
+            auto read = readRangesMatch->Read.Cast<TKqlReadTableRangesBase>();
+            if (TCoVoid::Match(read.Ranges().Raw())) {
+                auto key = Build<TKqlKeyInc>(ctx, read.Pos()).Done();
+                readMatch = readRangesMatch;
+                readMatch->Read =
+                    Build<TKqlReadTable>(ctx, read.Pos())
+                        .Settings(read.Settings())
+                        .Table(read.Table())
+                        .Columns(read.Columns())
+                        .Range<TKqlKeyRange>()
+                            .From(key)
+                            .To(key)
+                            .Build()
+                        .Done();
+                onlyPointRanges = true;
+            } else {
+                return node;
+            }
+        } else {
+            return node;
+        }
+    }
+
     if (!readMatch) {
         return node;
     }
@@ -159,7 +200,9 @@ TExprBase KqpPushPredicateToReadTable(TExprBase node, TExprContext& ctx, const T
 
     auto& tableDesc = indexName ? kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, mainTableDesc.Metadata->GetIndexMetadata(TString(indexName.Cast())).first->Name) : mainTableDesc;
 
-    YQL_ENSURE(tableDesc.Metadata->Kind != EKikimrTableKind::Olap);
+    if (tableDesc.Metadata->Kind == EKikimrTableKind::Olap) {
+        return node;
+    }
 
     auto row = flatmap.Lambda().Args().Arg(0);
     auto predicate = TExprBase(flatmap.Lambda().Body().Ref().ChildPtr(0));
@@ -178,6 +221,9 @@ TExprBase KqpPushPredicateToReadTable(TExprBase node, TExprContext& ctx, const T
 
     for (auto& keyRange : lookup.GetKeyRanges()) {
         bool useLookup = false;
+        if (onlyPointRanges && !IsPointPrefix(keyRange)) {
+            return node;
+        }
         if (keyRange.IsEquiRange()) {
             bool isFullKey = keyRange.GetNumDefined() == tableDesc.Metadata->KeyColumnNames.size();
 
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
index ea9b7825f0..118442443d 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_sqlin.cpp
@@ -42,6 +42,31 @@ TExprBase KqpRewriteSqlInToEquiJoin(const TExprBase& node, TExprContext& ctx, co
     }
 
     auto readMatch = MatchRead<TKqlReadTableBase>(flatMap.Input());
+
+    //TODO: remove this branch KIKIMR-15255
+    if (!readMatch) {
+        if (auto readRangesMatch = MatchRead<TKqlReadTableRangesBase>(flatMap.Input())) {
+            auto read = readRangesMatch->Read.Cast<TKqlReadTableRangesBase>();
+            if (TCoVoid::Match(read.Ranges().Raw())) {
+                readMatch = readRangesMatch;
+                auto key = Build<TKqlKeyInc>(ctx, read.Pos()).Done();
+                readMatch->Read =
+                    Build<TKqlReadTable>(ctx, read.Pos())
+                        .Settings(read.Settings())
+                        .Table(read.Table())
+                        .Columns(read.Columns())
+                        .Range<TKqlKeyRange>()
+                            .From(key)
+                            .To(key)
+                            .Build()
+                        .Done();
+            } else {
+                return node;
+            }
+        } else {
+            return node;
+        }
+    }
     if (!readMatch) {
         return node;
     }
diff --git a/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp b/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
index e9d6133b51..cec70f8353 100644
--- a/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_ne_flowcontrol_ut.cpp
@@ -57,7 +57,14 @@ void DoFlowControlTest(ui64 limit, bool hasBlockedByCapacity, bool useSessionAct
     appCfg.MutableTableServiceConfig()->MutableResourceManager()->SetMkqlHeavyProgramMemoryLimit(200ul << 20);
     appCfg.MutableTableServiceConfig()->MutableResourceManager()->SetQueryMemoryLimit(20ul << 30);
 
-    auto kikimr = KikimrRunnerEnableSessionActor(useSessionActor, {}, appCfg);
+    // TODO: KIKIMR-14294
+    auto kikimrSettings = TKikimrSettings()
+        .SetAppConfig(appCfg)
+        .SetEnableKqpSessionActor(useSessionActor)
+        .SetKqpSettings({})
+        .SetEnableKqpScanQueryStreamLookup(false);
+    TKikimrRunner kikimr{kikimrSettings};
+
     CreateSampleTables(kikimr);
     auto db = kikimr.GetTableClient();