YQL-13710 find less&equal procs from typeclass

ref:f6cb6dd9507a2a1ac87877c123d670ad7d2c349d

4 files changed, 149 insertions, 3 deletions

diff --git a/ydb/library/yql/parser/pg_catalog/CMakeLists.txt b/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
index e3ca4f495bf..27f5f29ae45 100644
--- a/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
+++ b/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
@@ -31,6 +31,7 @@ target_sources(yql-parser-pg_catalog.global PRIVATE
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/15e039bc35a86f84476091e328dd74ea.cpp
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/ca7bb755fe32f897eb03cf6004d1ad2a.cpp
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/f2374a475345069a2bfd38b387dfa11a.cpp
+  ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/462597d51509ecdc14b417ce082dc53b.cpp
 )
 resources(yql-parser-pg_catalog.global
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/99bdd9e510834dd355c2457fcccc53d5.cpp
@@ -81,3 +82,10 @@ resources(yql-parser-pg_catalog.global
   KEYS
   pg_amproc.dat
 )
+resources(yql-parser-pg_catalog.global
+  ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/462597d51509ecdc14b417ce082dc53b.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/contrib/libs/postgresql/src/include/catalog/pg_amop.dat
+  KEYS
+  pg_amop.dat
+)
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.cpp b/ydb/library/yql/parser/pg_catalog/catalog.cpp
index b9e4bb2774a..b8e51630ffa 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.cpp
+++ b/ydb/library/yql/parser/pg_catalog/catalog.cpp
@@ -34,8 +34,11 @@ using TAggregations = THashMap<ui32, TAggregateDesc>;
 
 using TOpClasses = THashMap<std::pair<EOpClassMethod, ui32>, TOpClassDesc>;
 
+using TAmOps = THashMap<std::tuple<EOpClassMethod, TString, ui32, ui32, ui32>, TAmOpDesc>;
+
 class TParser {
 public:
+    virtual ~TParser() = default;
     void Do(const TString& dat) {
         enum class EState {
             WaitBracket,
@@ -635,7 +638,7 @@ public:
     void OnFinish() override {
         if (IsSupported) {
             Y_ENSURE(!LastOpClass.Name.empty());
-            OpClasses[std::pair(LastOpClass.Method, LastOpClass.TypeId)] = LastOpClass;
+            OpClasses[std::make_pair(LastOpClass.Method, LastOpClass.TypeId)] = LastOpClass;
         }
 
         IsSupported = true;
@@ -649,6 +652,74 @@ private:
     bool IsSupported = true;
 };
 
+class TAmOpsParser : public TParser {
+public:
+    TAmOpsParser(TAmOps& amOps, const THashMap<TString, ui32>& typeByName,
+        const THashMap<TString, TVector<ui32>>& operatorsByName, const TOperators& operators)
+        : AmOps(amOps)
+        , TypeByName(typeByName)
+        , OperatorsByName(operatorsByName)
+        , Operators(operators)
+    {}
+
+    void OnKey(const TString& key, const TString& value) override {
+        if (key == "amopmethod") {
+            if (value == "btree") {
+                LastAmOp.Method = EOpClassMethod::Btree;
+            } else {
+                IsSupported = false;
+            }
+        } else if (key == "amopfamily") {
+            LastAmOp.Family = value;
+        } else if (key == "amoplefttype") {
+            auto leftTypePtr = TypeByName.FindPtr(value);
+            Y_ENSURE(leftTypePtr);
+            LastAmOp.LeftType = *leftTypePtr;
+        } else if (key == "amoprighttype") {
+            auto rightTypePtr = TypeByName.FindPtr(value);
+            Y_ENSURE(rightTypePtr);
+            LastAmOp.RightType = *rightTypePtr;
+        } else if (key == "amopstrategy") {
+            LastAmOp.Strategy = FromString<ui32>(value);
+        } else if (key == "amopopr") {
+            auto pos = value.find('(');
+            Y_ENSURE(pos != TString::npos);
+            LastOp = value.substr(0, pos);
+        }
+    }
+
+    void OnFinish() override {
+        if (IsSupported) {
+            auto operIdPtr = OperatorsByName.FindPtr(LastOp);
+            Y_ENSURE(operIdPtr);
+            for (const auto& id : *operIdPtr) {
+                const auto& d = Operators.FindPtr(id);
+                Y_ENSURE(d);
+                if (d->Kind == EOperKind::Binary && IsCompatibleTo(LastAmOp.LeftType, d->LeftType) && IsCompatibleTo(LastAmOp.RightType, d->RightType)) {
+                    Y_ENSURE(!LastAmOp.OperId);
+                    LastAmOp.OperId = d->OperId;
+                }
+            }
+
+            Y_ENSURE(LastAmOp.OperId);
+            AmOps[std::make_tuple(LastAmOp.Method, LastAmOp.Family, LastAmOp.Strategy, LastAmOp.LeftType, LastAmOp.RightType)] = LastAmOp;
+        }
+
+        IsSupported = true;
+        LastAmOp = TAmOpDesc();
+        LastOp = "";
+    }
+
+private:
+    TAmOps& AmOps;
+    const THashMap<TString, ui32>& TypeByName;
+    const THashMap<TString, TVector<ui32>>& OperatorsByName;
+    const TOperators& Operators;
+    TAmOpDesc LastAmOp;
+    bool IsSupported = true;
+    TString LastOp;
+};
+
 TOperators ParseOperators(const TString& dat, const THashMap<TString, ui32>& typeByName,
     const THashMap<TString, TVector<ui32>>& procByName) {
     TOperators ret;
@@ -694,6 +765,14 @@ TOpClasses ParseOpClasses(const TString& dat, const THashMap<TString, ui32>& typ
     return ret;
 }
 
+TAmOps ParseAmOps(const TString& dat, const THashMap<TString, ui32>& typeByName,
+    const THashMap<TString, TVector<ui32>>& operatorsByName, const TOperators& operators) {
+    TAmOps ret;
+    TAmOpsParser parser(ret, typeByName, operatorsByName, operators);
+    parser.Do(dat);
+    return ret;
+}
+
 struct TCatalog {
     TCatalog() {
         TString typeData;
@@ -710,6 +789,8 @@ struct TCatalog {
         Y_ENSURE(NResource::FindExact("pg_opclass.dat", &opClassData));
         TString amProcData;
         Y_ENSURE(NResource::FindExact("pg_amproc.dat", &amProcData));
+        TString amOpData;
+        Y_ENSURE(NResource::FindExact("pg_amop.dat", &amOpData));
         THashMap<ui32, TLazyTypeInfo> lazyTypeInfos;
         Types = ParseTypes(typeData, lazyTypeInfos);
         for (const auto& [k, v] : Types) {
@@ -804,6 +885,24 @@ struct TCatalog {
         }
 
         OpClasses = ParseOpClasses(opClassData, TypeByName);
+        AmOps = ParseAmOps(amOpData, TypeByName, OperatorsByName, Operators);
+        for (auto&[k, v] : Types) {
+            if (v.TypeId != v.ArrayTypeId) {
+                auto opClassPtr = OpClasses.FindPtr(std::make_pair(EOpClassMethod::Btree, v.TypeId));
+                if (opClassPtr) {
+                    auto lessAmOpPtr = AmOps.FindPtr(std::make_tuple(EOpClassMethod::Btree, opClassPtr->Family, ui32(EBtreeAmStrategy::Less), v.TypeId, v.TypeId));
+                    Y_ENSURE(lessAmOpPtr);
+                    auto equalAmOpPtr = AmOps.FindPtr(std::make_tuple(EOpClassMethod::Btree, opClassPtr->Family, ui32(EBtreeAmStrategy::Equal), v.TypeId, v.TypeId));
+                    Y_ENSURE(equalAmOpPtr);
+                    auto lessOperPtr = Operators.FindPtr(lessAmOpPtr->OperId);
+                    Y_ENSURE(lessOperPtr);
+                    auto equalOperPtr = Operators.FindPtr(equalAmOpPtr->OperId);
+                    Y_ENSURE(equalOperPtr);
+                    v.LessProcId = lessOperPtr->ProcId;
+                    v.EqualProcId = equalOperPtr->ProcId;
+                }
+            }
+        }
     }
 
     static const TCatalog& Instance() {
@@ -816,6 +915,7 @@ struct TCatalog {
     TCasts Casts;
     TAggregations Aggregations;
     TOpClasses OpClasses;
+    TAmOps AmOps;
     THashMap<TString, TVector<ui32>> ProcByName;
     THashMap<TString, ui32> TypeByName;
     THashMap<std::pair<ui32, ui32>, ui32> CastsByDir;
@@ -1026,12 +1126,12 @@ const TAggregateDesc& LookupAggregation(const TStringBuf& name, const TVector<ui
 
 bool HasOpClass(EOpClassMethod method, ui32 typeId) {
     const auto& catalog = TCatalog::Instance();
-    return catalog.OpClasses.contains(std::pair(method, typeId));
+    return catalog.OpClasses.contains(std::make_pair(method, typeId));
 }
 
 const TOpClassDesc& LookupOpClass(EOpClassMethod method, ui32 typeId) {
     const auto& catalog = TCatalog::Instance();
-    auto opClassPtr = catalog.OpClasses.FindPtr(std::pair(method, typeId));
+    auto opClassPtr = catalog.OpClasses.FindPtr(std::make_pair(method, typeId));
     if (!opClassPtr) {
         throw yexception() << "No such opclass";
     }
@@ -1039,4 +1139,16 @@ const TOpClassDesc& LookupOpClass(EOpClassMethod method, ui32 typeId) {
     return *opClassPtr;
 }
 
+/*
+const TAmOpDesc& LookupAmOp(EOpClassMethod method, const TString& family, ui32 strategy, ui32 leftType, ui32 rightType) {
+    const auto& catalog = TCatalog::Instance();
+    auto amOpPtr = catalog.AmOps.FindPtr(std::make_tuple(method, family, strategy, leftType, rightType));
+    if (!amOpPtr) {
+        throw yexception() << "No such amop";
+    }
+
+    return *amOpPtr;
+}
+*/
+
 }
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.h b/ydb/library/yql/parser/pg_catalog/catalog.h
index 778442aeddf..05feac2c993 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.h
+++ b/ydb/library/yql/parser/pg_catalog/catalog.h
@@ -42,6 +42,9 @@ struct TTypeDesc {
     ui32 SendFuncId = 0;
     ui32 ReceiveFuncId = 0;
     i32 TypeLen = 0;
+    // from opclass
+    ui32 LessProcId = 0;
+    ui32 EqualProcId = 0;
 };
 
 enum class ECastMethod {
@@ -88,6 +91,23 @@ struct TOpClassDesc {
     TString Family;
 };
 
+struct TAmOpDesc {
+    EOpClassMethod Method = EOpClassMethod::Btree;
+    TString Family;
+    ui32 Strategy = 0;
+    ui32 LeftType = 0;
+    ui32 RightType = 0;
+    ui32 OperId = 0;
+};
+
+enum class EBtreeAmStrategy {
+    Less = 1,
+    LessOrEqual = 2,
+    Equal = 3,
+    GreaterOrEqual = 4,
+    Greater = 5
+};
+
 const TProcDesc& LookupProc(const TString& name, const TVector<ui32>& argTypeIds);
 const TProcDesc& LookupProc(ui32 procId, const TVector<ui32>& argTypeIds);
 const TProcDesc& LookupProc(ui32 procId);
@@ -109,6 +129,8 @@ const TAggregateDesc& LookupAggregation(const TStringBuf& name, const TVector<ui
 bool HasOpClass(EOpClassMethod method, ui32 typeId);
 const TOpClassDesc& LookupOpClass(EOpClassMethod method, ui32 typeId);
 
+const TAmOpDesc& LookupAmOp(EOpClassMethod method, const TString& family, ui32 strategy, ui32 leftType, ui32 rightType);
+
 bool IsCompatibleTo(ui32 actualType, ui32 expectedType);
 
 }
diff --git a/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp b/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp
index 91b4395791a..af818bc7ebd 100644
--- a/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp
+++ b/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp
@@ -26,12 +26,16 @@ Y_UNIT_TEST_SUITE(TTypesTests) {
         UNIT_ASSERT_VALUES_EQUAL(ret.ArrayTypeId, 1009);
         UNIT_ASSERT_VALUES_EQUAL(ret.Name, "text");
         UNIT_ASSERT_VALUES_EQUAL(ret.ElementTypeId, 0);
+        UNIT_ASSERT(ret.LessProcId);
+        UNIT_ASSERT(ret.EqualProcId);
 
         ret = LookupType("point");
         UNIT_ASSERT_VALUES_EQUAL(ret.TypeId, 600);
         UNIT_ASSERT_VALUES_EQUAL(ret.ArrayTypeId, 1017);
         UNIT_ASSERT_VALUES_EQUAL(ret.Name, "point");
         UNIT_ASSERT_VALUES_EQUAL(ret.ElementTypeId, LookupType("float8").TypeId);
+        UNIT_ASSERT(!ret.LessProcId);
+        UNIT_ASSERT(!ret.EqualProcId);
 
         ret = LookupType(1009);
         UNIT_ASSERT_VALUES_EQUAL(ret.TypeId, 1009);