KIKIMR-19719: Lookup for simple PgSelect

ydb/core/kqp/opt/kqp_query_plan.cpp

@@ -707,8 +707,8 @@ class TxPlanSerializer {
                             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);
@@ -1463,8 +1463,8 @@ class TxPlanSerializer {
                         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);

ydb/core/kqp/opt/logical/kqp_opt_log.cpp

@@ -28,8 +28,8 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
         , 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));

ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp

@@ -179,11 +179,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;
@@ -250,6 +250,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();
         }
@@ -498,10 +500,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

ydb/core/kqp/opt/query_plan_value/kqp_query_plan_value.cpp

@@ -18,5 +18,15 @@ TString TValue::GetPgText() const {
     return convertResult.Str;
 }
 
+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()));
+}
+
 }
 }

ydb/core/kqp/ut/pg/kqp_pg_ut.cpp

@@ -3396,7 +3396,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\""));
             }
         }
 
@@ -3475,6 +3476,73 @@ 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)));
+        }
+   }
 }
 
 } // namespace NKqp

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>()) {

ydb/public/lib/value/value.h

@@ -89,9 +89,13 @@ class TValue {
     NScheme::TTypeId GetDataType() const;
     // gets text representation of simple 'Data' types
     TString GetDataText() const;
+
+    // You need to add ydb/core/kqp/opt/query_plan_value to PEERDIRs in order to use all Pg-related functions
     // 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