KIKIMR-19719: Lookup for simple PgSelect (#839)

10 files changed, 177 insertions, 49 deletions

diff --git a/ydb/core/kqp/opt/kqp_query_plan.cpp b/ydb/core/kqp/opt/kqp_query_plan.cpp
index 26a791b4493..d2778ba244e 100644
--- a/ydb/core/kqp/opt/kqp_query_plan.cpp
+++ b/ydb/core/kqp/opt/kqp_query_plan.cpp
@@ -707,8 +707,8 @@ private:
                             TStringBuilder rangeDesc;
                             rangeDesc << keyColumns[colId] << " "
                                 << (from[keyColumns.size()].GetDataText() == "1" ? "[" : "(")
-                                << (from[colId].HaveValue() ? from[colId].GetDataText() : "-∞") << ", "
-                                << (to[colId].HaveValue() ? to[colId].GetDataText() : "+∞")
+                                << (from[colId].HaveValue() ? from[colId].GetSimpleValueText() : "-∞") << ", "
+                                << (to[colId].HaveValue() ? to[colId].GetSimpleValueText() : "+∞")
                                 << (to[keyColumns.size()].GetDataText() == "1" ? "]" : ")");
 
                             readInfo.ScanBy.push_back(rangeDesc);
@@ -1469,8 +1469,8 @@ private:
                         TStringBuilder rangeDesc;
                         rangeDesc << keyColumns[colId] << " "
                             << (from[keyColumns.size()].GetDataText() == "1" ? "[" : "(")
-                            << (from[colId].HaveValue() ? from[colId].GetDataText() : "-∞") << ", "
-                            << (to[colId].HaveValue() ? to[colId].GetDataText() : "+∞")
+                            << (from[colId].HaveValue() ? from[colId].GetSimpleValueText() : "-∞") << ", "
+                            << (to[colId].HaveValue() ? to[colId].GetSimpleValueText() : "+∞")
                             << (to[keyColumns.size()].GetDataText() == "1" ? "]" : ")");
 
                         readInfo.ScanBy.push_back(rangeDesc);
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
index 36efc2906c5..96f7055653f 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -28,8 +28,8 @@ public:
         , Config(config)
     {
 #define HNDL(name) "KqpLogical-"#name, Hndl(&TKqpLogicalOptTransformer::name)
-        AddHandler(0, &TCoFlatMap::Match, HNDL(PushPredicateToReadTable));
-        AddHandler(0, &TCoFlatMap::Match, HNDL(PushExtractedPredicateToReadTable));
+        AddHandler(0, &TCoFlatMapBase::Match, HNDL(PushPredicateToReadTable));
+        AddHandler(0, &TCoFlatMapBase::Match, HNDL(PushExtractedPredicateToReadTable));
         AddHandler(0, &TCoAggregate::Match, HNDL(RewriteAggregate));
         AddHandler(0, &TCoAggregateCombine::Match, HNDL(PushdownOlapGroupByKeys));
         AddHandler(0, &TCoTake::Match, HNDL(RewriteTakeSortToTopSort));
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
index b3aeaed1978..5a8ff001bcd 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
@@ -181,11 +181,11 @@ TMaybeNode<TExprBase> TryBuildTrivialReadTable(TCoFlatMap& flatmap, TKqlReadTabl
 TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx,
     TTypeAnnotationContext& typesCtx)
 {
-    if (!node.Maybe<TCoFlatMap>()) {
+    if (!node.Maybe<TCoFlatMapBase>()) {
         return node;
     }
 
-    auto flatmap = node.Cast<TCoFlatMap>();
+    auto flatmap = node.Cast<TCoFlatMapBase>();
 
     if (!IsPredicateFlatMap(flatmap.Lambda().Body().Ref())) {
         return node;
@@ -252,6 +252,8 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
                 kqpCtx.Cluster,
                 mainTableDesc.Metadata->GetIndexMetadata(TString(indexName.Cast())).first->Name)
             : mainTableDesc;
+        YQL_ENSURE(node.Maybe<TCoFlatMap>(), "got OrderedFlatMap with disabled PredicateExtract20");
+        auto flatmap = node.Cast<TCoFlatMap>();
         if (auto expr = TryBuildTrivialReadTable(flatmap, read, *readMatch, tableDesc, ctx, kqpCtx, indexName)) {
             return expr.Cast();
         }
@@ -501,10 +503,17 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
 
     *input = readMatch->BuildProcessNodes(*input, ctx);
 
-    return Build<TCoFlatMap>(ctx, node.Pos())
-        .Input(*input)
-        .Lambda(residualLambda)
-        .Done();
+    if (node.Maybe<TCoFlatMap>()) {
+        return Build<TCoFlatMap>(ctx, node.Pos())
+            .Input(*input)
+            .Lambda(residualLambda)
+            .Done();
+    } else {
+        return Build<TCoOrderedFlatMap>(ctx, node.Pos())
+            .Input(*input)
+            .Lambda(residualLambda)
+            .Done();    
+    }
 }
 
 } // namespace NKikimr::NKqp::NOpt
diff --git a/ydb/core/kqp/opt/query_plan_value/kqp_query_plan_value.cpp b/ydb/core/kqp/opt/query_plan_value/kqp_query_plan_value.cpp
deleted file mode 100644
index b72f09d47de..00000000000
--- a/ydb/core/kqp/opt/query_plan_value/kqp_query_plan_value.cpp
+++ /dev/null
@@ -1,22 +0,0 @@
-#include <ydb/library/yql/parser/pg_wrapper/interface/type_desc.h>
-#include <ydb/public/lib/value/value.h>
-
-namespace NKikimr {
-namespace NClient {
-
-TString TValue::GetPgText() const {
-    Y_ASSERT(Type.GetKind() == NKikimrMiniKQL::ETypeKind::Pg);
-    if (Value.HasNullFlagValue()) {
-        return TString("null");
-    }
-    if (Value.HasText()) {
-        return Value.GetText();
-    }
-    auto pgType = Type.GetPg();
-    auto convertResult = NPg::PgNativeTextFromNativeBinary(Value.GetBytes(), NPg::TypeDescFromPgTypeId(pgType.Getoid()));
-    Y_ENSURE(!convertResult.Error, convertResult.Error);
-    return convertResult.Str;
-}
-
-}
-}
diff --git a/ydb/core/kqp/opt/query_plan_value/ya.make b/ydb/core/kqp/opt/query_plan_value/ya.make
deleted file mode 100644
index 1424b3fafdc..00000000000
--- a/ydb/core/kqp/opt/query_plan_value/ya.make
+++ /dev/null
@@ -1,12 +0,0 @@
-LIBRARY()
-
-SRCS(
-    kqp_query_plan_value.cpp
-)
-
-PEERDIR(
-    ydb/library/yql/parser/pg_wrapper/interface
-    ydb/public/lib/value
-)
-
-END()
diff --git a/ydb/core/kqp/opt/ya.make b/ydb/core/kqp/opt/ya.make
index 6de6e79023c..816583d41eb 100644
--- a/ydb/core/kqp/opt/ya.make
+++ b/ydb/core/kqp/opt/ya.make
@@ -21,7 +21,6 @@ PEERDIR(
     ydb/core/kqp/opt/logical
     ydb/core/kqp/opt/peephole
     ydb/core/kqp/opt/physical
-    ydb/core/kqp/opt/query_plan_value
     ydb/library/yql/dq/common
     ydb/library/yql/dq/opt
     ydb/library/yql/dq/type_ann
diff --git a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
index 5f32e98be20..2f456bb9e7f 100644
--- a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
+++ b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
@@ -3529,7 +3529,8 @@ Y_UNIT_TEST_SUITE(KqpPg) {
                     SELECT * FROM PgTable WHERE key = 'a';
                 )");
                 auto result = db.ExecuteQuery(query, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
-                UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+                UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+                UNIT_ASSERT(result.GetIssues().ToString().Contains("invalid input syntax for type integer: \"a\""));
             }
         }
 
@@ -3608,6 +3609,127 @@ Y_UNIT_TEST_SUITE(KqpPg) {
             UNIT_ASSERT(result.GetIssues().ToString().Contains("invalid byte sequence for encoding \"UTF8\": 0x00"));
         }
     }
+
+    Y_UNIT_TEST(NoSelectFullScan) {
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetEnablePreparedDdl(true);
+        auto setting = NKikimrKqp::TKqpSetting();
+        auto serverSettings = TKikimrSettings()
+            .SetAppConfig(appConfig)
+            .SetKqpSettings({setting});
+        TKikimrRunner kikimr(serverSettings.SetWithSampleTables(false));
+        auto db = kikimr.GetQueryClient();
+        auto settings = NYdb::NQuery::TExecuteQuerySettings().Syntax(NYdb::NQuery::ESyntax::Pg);
+        {
+            auto result = db.ExecuteQuery(R"(
+                CREATE TABLE pgbench_accounts(aid    int not null,bid int,abalance int,filler char(84), primary key (aid))
+            )", NYdb::NQuery::TTxControl::NoTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                INSERT INTO pgbench_accounts (aid, bid, abalance, filler) VALUES
+                    (1, 1, 10, '                                                                                    '::char),
+                    (2, 1, 20, '                                                                                    '::char),
+                    (3, 1, 30, '                                                                                    '::char),
+                    (4, 1, 40, '
+                                               '::char),
+                    (5, 1, 50, '                                                                                    '::char),
+                    (6, 1, 60, '                                                                                    '::char),
+                    (7, 1, 70, '                                                                                    '::char),
+                    (8, 1, 80, '                                                                                    '::char),
+                    (9, 1, 90, '                                                                                    '::char),
+                    (10, 1, 100, '                                                                                    '::char)
+            )", NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+        {
+            auto tc = kikimr.GetTableClient();
+            TStreamExecScanQuerySettings settings;
+            settings.Explain(true);
+            auto it = tc.StreamExecuteScanQuery(R"(
+                --!syntax_pg
+                SELECT abalance FROM pgbench_accounts WHERE aid = 7 OR aid = 3 ORDER BY abalance;
+            )", settings).GetValueSync();
+
+            UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
+
+            auto res = CollectStreamResult(it);
+            UNIT_ASSERT(res.PlanJson);
+
+            NJson::TJsonValue plan;
+            NJson::ReadJsonTree(*res.PlanJson, &plan, true);
+            UNIT_ASSERT(ValidatePlanNodeIds(plan));
+
+            auto fullScan = FindPlanNodeByKv(plan, "Node Type", "Filter-TableFullScan");
+            UNIT_ASSERT_C(!fullScan.IsDefined(), "got fullscan, expected lookup");
+            auto lookup = FindPlanNodeByKv(plan, "Node Type", "TableLookup");
+            UNIT_ASSERT_C(lookup.IsDefined(), "no Table Lookup in plan");
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                SELECT abalance FROM pgbench_accounts WHERE aid = 7 OR aid = 3 ORDER BY abalance;
+            )", NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            CompareYson(R"([
+                ["30"];["70"]
+            ])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+        {
+            auto tc = kikimr.GetTableClient();
+            TStreamExecScanQuerySettings settings;
+            settings.Explain(true);
+            auto it = tc.StreamExecuteScanQuery(R"(
+                --!syntax_pg
+                SELECT abalance FROM pgbench_accounts WHERE aid = 7 OR aid < 3 ORDER BY abalance;
+            )", settings).GetValueSync();
+
+            UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
+
+            auto res = CollectStreamResult(it);
+            UNIT_ASSERT(res.PlanJson);
+            Cerr << res.PlanJson << Endl;
+
+            NJson::TJsonValue plan;
+            NJson::ReadJsonTree(*res.PlanJson, &plan, true);
+            UNIT_ASSERT(ValidatePlanNodeIds(plan));
+
+            auto fullScan = FindPlanNodeByKv(plan, "Node Type", "Filter-TableFullScan");
+            UNIT_ASSERT_C(!fullScan.IsDefined(), "got fullscan, expected lookup");
+            auto lookup = FindPlanNodeByKv(plan, "Node Type", "TableRangeScan");
+            UNIT_ASSERT_C(lookup.IsDefined(), "no Table Range Scan in plan");
+        }
+        {
+            auto tc = kikimr.GetTableClient();
+            TStreamExecScanQuerySettings settings;
+            settings.Explain(true);
+            auto it = tc.StreamExecuteScanQuery(R"(
+                --!syntax_pg
+                SELECT abalance FROM pgbench_accounts WHERE aid > 4 AND aid < 3;
+            )", settings).GetValueSync();
+
+            UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
+
+            auto res = CollectStreamResult(it);
+            UNIT_ASSERT(res.PlanJson);
+            Cerr << res.PlanJson << Endl;
+            NJson::TJsonValue plan;
+            NJson::ReadJsonTree(*res.PlanJson, &plan, true);
+            UNIT_ASSERT(ValidatePlanNodeIds(plan));
+
+            auto fullScan = FindPlanNodeByKv(plan, "Node Type", "Filter-TableFullScan");
+            UNIT_ASSERT_C(!fullScan.IsDefined(), "got fullscan, expected lookup");
+            auto lookup = FindPlanNodeByKv(plan, "Node Type", "TableRangeScan");
+            UNIT_ASSERT_C(lookup.IsDefined(), "no Table Range Scan in plan");
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                SELECT abalance FROM pgbench_accounts WHERE aid > 4 AND aid < 3;
+            )", NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            CompareYson(R"([])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+   }
 }
 
 } // namespace NKqp
diff --git a/ydb/library/yql/utils/plan/plan_utils.cpp b/ydb/library/yql/utils/plan/plan_utils.cpp
index 00454ff3957..2a1e8e32bd2 100644
--- a/ydb/library/yql/utils/plan/plan_utils.cpp
+++ b/ydb/library/yql/utils/plan/plan_utils.cpp
@@ -27,6 +27,13 @@ TString ToStr(const TCoDataCtor& data) {
     return out.Str();
 }
 
+TString ToStr(const TCoPgConst& data) {
+    TStringStream out;
+    EscapeArbitraryAtom(data.Value().Value(), '"', &out);
+    return out.Str();
+}
+
+
 TString ToStr(const TCoLambda& lambda) {
     return PrettyExprStr(lambda.Body());
 }
@@ -143,6 +150,8 @@ TString PrettyExprStr(const TExprBase& expr) {
         return TString(expr.Ref().Child(0)->Content());
     } else if (auto data = expr.Maybe<TCoDataCtor>()) {
         return ToStr(data.Cast());
+    } else if (auto pgConst = expr.Maybe<TCoPgConst>()) {
+        return ToStr(pgConst.Cast());
     } else if (auto lambda = expr.Maybe<TCoLambda>()) {
         return ToStr(lambda.Cast());
     } else if (auto asStruct = expr.Maybe<TCoAsStruct>()) {
diff --git a/ydb/public/lib/value/value.cpp b/ydb/public/lib/value/value.cpp
index 7edde8c7a5e..f96d235c68d 100644
--- a/ydb/public/lib/value/value.cpp
+++ b/ydb/public/lib/value/value.cpp
@@ -432,6 +432,26 @@ TString TValue::GetDataText() const {
     return TStringBuilder() << "\"<unknown type "  << Type.GetData().GetScheme() << ">\"";
 }
 
+TString TValue::GetPgText() const {
+    Y_ASSERT(Type.GetKind() == NKikimrMiniKQL::ETypeKind::Pg);
+    if (Value.HasNullFlagValue()) {
+        return TString("null");
+    }
+    Y_ENSURE(Value.HasText());
+    return Value.GetText();
+}
+
+TString TValue::GetSimpleValueText() const {
+    if (Type.GetKind() == NKikimrMiniKQL::ETypeKind::Pg) {
+        return GetPgText();
+    }
+    if (Type.GetKind() == NKikimrMiniKQL::ETypeKind::Data) {
+        return GetDataText();
+    }
+    Y_ENSURE(false, TStringBuilder() << "unexpected NKikimrMiniKQL::ETypeKind: " << ETypeKind_Name(GetType().GetKind()));
+}
+
+
 template <> TString TValue::GetTypeText<TFormatCxx>(const TFormatCxx& format) const {
     switch(Type.GetKind()) {
     case NKikimrMiniKQL::ETypeKind::Void:
diff --git a/ydb/public/lib/value/value.h b/ydb/public/lib/value/value.h
index f72cccfc232..024cbbbb52c 100644
--- a/ydb/public/lib/value/value.h
+++ b/ydb/public/lib/value/value.h
@@ -89,9 +89,12 @@ public:
     NScheme::TTypeId GetDataType() const;
     // gets text representation of simple 'Data' types
     TString GetDataText() const;
+
     // gets text representation of simple 'Pg' types
-    // You need to add ydb/core/kqp/opt/query_plan_value to PEERDIRs in order to use this function
     TString GetPgText() const;
+    // gets text representation of simple 'Data' and 'Pg' types
+    TString GetSimpleValueText() const;
+
     // returns text representation of value's type
     template <typename Format> TString GetTypeText(const Format& format = Format()) const;
     // returns text representation of value itself