YQL-13710 parse opclasses

ref:c3f6a9f4e3763db0d9240b83c7ec32ef4d185045

4 files changed, 126 insertions, 0 deletions

diff --git a/ydb/library/yql/parser/pg_catalog/CMakeLists.txt b/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
index 768259591bf..e3ca4f495bf 100644
--- a/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
+++ b/ydb/library/yql/parser/pg_catalog/CMakeLists.txt
@@ -29,6 +29,8 @@ target_sources(yql-parser-pg_catalog.global PRIVATE
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/1ecd65c896f36a3990d870644a6da9c8.cpp
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/46b25697572e7e60703079b71cd18295.cpp
   ${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
 )
 resources(yql-parser-pg_catalog.global
   ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/99bdd9e510834dd355c2457fcccc53d5.cpp
@@ -65,3 +67,17 @@ resources(yql-parser-pg_catalog.global
   KEYS
   pg_aggregate.dat
 )
+resources(yql-parser-pg_catalog.global
+  ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/ca7bb755fe32f897eb03cf6004d1ad2a.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/contrib/libs/postgresql/src/include/catalog/pg_opclass.dat
+  KEYS
+  pg_opclass.dat
+)
+resources(yql-parser-pg_catalog.global
+  ${CMAKE_BINARY_DIR}/ydb/library/yql/parser/pg_catalog/f2374a475345069a2bfd38b387dfa11a.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/contrib/libs/postgresql/src/include/catalog/pg_amproc.dat
+  KEYS
+  pg_amproc.dat
+)
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.cpp b/ydb/library/yql/parser/pg_catalog/catalog.cpp
index 1dc8aa8d566..b9e4bb2774a 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.cpp
+++ b/ydb/library/yql/parser/pg_catalog/catalog.cpp
@@ -32,6 +32,8 @@ using TCasts = THashMap<ui32, TCastDesc>;
 
 using TAggregations = THashMap<ui32, TAggregateDesc>;
 
+using TOpClasses = THashMap<std::pair<EOpClassMethod, ui32>, TOpClassDesc>;
+
 class TParser {
 public:
     void Do(const TString& dat) {
@@ -603,6 +605,50 @@ private:
     TString LastDeserializeFunc;
 };
 
+class TOpClassesParser : public TParser {
+public:
+    TOpClassesParser(TOpClasses& opClasses, const THashMap<TString, ui32>& typeByName)
+        : OpClasses(opClasses)
+        , TypeByName(typeByName)
+    {}
+
+    void OnKey(const TString& key, const TString& value) override {
+        if (key == "opcmethod") {
+            if (value == "btree") {
+                LastOpClass.Method = EOpClassMethod::Btree;
+            } else if (value == "hash") {
+                LastOpClass.Method = EOpClassMethod::Hash;
+            } else {
+                IsSupported = false;
+            }
+        } else if (key == "opcintype") {
+            auto idPtr = TypeByName.FindPtr(value);
+            Y_ENSURE(idPtr);
+            LastOpClass.TypeId = *idPtr;
+        } else if (key == "opcname") {
+            LastOpClass.Name = value;
+        } else if (key == "opcfamily") {
+            LastOpClass.Family = value;
+        }
+    }
+
+    void OnFinish() override {
+        if (IsSupported) {
+            Y_ENSURE(!LastOpClass.Name.empty());
+            OpClasses[std::pair(LastOpClass.Method, LastOpClass.TypeId)] = LastOpClass;
+        }
+
+        IsSupported = true;
+        LastOpClass = TOpClassDesc();
+    }
+
+private:
+    TOpClasses& OpClasses;
+    const THashMap<TString, ui32>& TypeByName;
+    TOpClassDesc LastOpClass;
+    bool IsSupported = true;
+};
+
 TOperators ParseOperators(const TString& dat, const THashMap<TString, ui32>& typeByName,
     const THashMap<TString, TVector<ui32>>& procByName) {
     TOperators ret;
@@ -641,6 +687,13 @@ TCasts ParseCasts(const TString& dat, const THashMap<TString, ui32>& typeByName,
     return ret;
 }
 
+TOpClasses ParseOpClasses(const TString& dat, const THashMap<TString, ui32>& typeByName) {
+    TOpClasses ret;
+    TOpClassesParser parser(ret, typeByName);
+    parser.Do(dat);
+    return ret;
+}
+
 struct TCatalog {
     TCatalog() {
         TString typeData;
@@ -653,6 +706,10 @@ struct TCatalog {
         Y_ENSURE(NResource::FindExact("pg_cast.dat", &castData));
         TString aggData;
         Y_ENSURE(NResource::FindExact("pg_aggregate.dat", &aggData));
+        TString opClassData;
+        Y_ENSURE(NResource::FindExact("pg_opclass.dat", &opClassData));
+        TString amProcData;
+        Y_ENSURE(NResource::FindExact("pg_amproc.dat", &amProcData));
         THashMap<ui32, TLazyTypeInfo> lazyTypeInfos;
         Types = ParseTypes(typeData, lazyTypeInfos);
         for (const auto& [k, v] : Types) {
@@ -745,6 +802,8 @@ struct TCatalog {
         for (const auto&[k, v] : Aggregations) {
             AggregationsByName[v.Name].push_back(k);
         }
+
+        OpClasses = ParseOpClasses(opClassData, TypeByName);
     }
 
     static const TCatalog& Instance() {
@@ -756,6 +815,7 @@ struct TCatalog {
     TTypes Types;
     TCasts Casts;
     TAggregations Aggregations;
+    TOpClasses OpClasses;
     THashMap<TString, TVector<ui32>> ProcByName;
     THashMap<TString, ui32> TypeByName;
     THashMap<std::pair<ui32, ui32>, ui32> CastsByDir;
@@ -964,4 +1024,19 @@ const TAggregateDesc& LookupAggregation(const TStringBuf& name, const TVector<ui
     throw yexception() << "Unable to find an overload for aggregate " << name << " with given argument types";
 }
 
+bool HasOpClass(EOpClassMethod method, ui32 typeId) {
+    const auto& catalog = TCatalog::Instance();
+    return catalog.OpClasses.contains(std::pair(method, typeId));
+}
+
+const TOpClassDesc& LookupOpClass(EOpClassMethod method, ui32 typeId) {
+    const auto& catalog = TCatalog::Instance();
+    auto opClassPtr = catalog.OpClasses.FindPtr(std::pair(method, typeId));
+    if (!opClassPtr) {
+        throw yexception() << "No such opclass";
+    }
+
+    return *opClassPtr;
+}
+
 }
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.h b/ydb/library/yql/parser/pg_catalog/catalog.h
index 4a67d4414d0..778442aeddf 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.h
+++ b/ydb/library/yql/parser/pg_catalog/catalog.h
@@ -76,6 +76,18 @@ struct TAggregateDesc {
     TString InitValue;
 };
 
+enum class EOpClassMethod {
+    Btree,
+    Hash
+};
+
+struct TOpClassDesc {
+    EOpClassMethod Method = EOpClassMethod::Btree;
+    ui32 TypeId = 0;
+    TString Name;
+    TString Family;
+};
+
 const TProcDesc& LookupProc(const TString& name, const TVector<ui32>& argTypeIds);
 const TProcDesc& LookupProc(ui32 procId, const TVector<ui32>& argTypeIds);
 const TProcDesc& LookupProc(ui32 procId);
@@ -94,6 +106,9 @@ const TOperDesc& LookupOper(ui32 operId);
 bool HasAggregation(const TStringBuf& name);
 const TAggregateDesc& LookupAggregation(const TStringBuf& name, const TVector<ui32>& argTypeIds);
 
+bool HasOpClass(EOpClassMethod method, ui32 typeId);
+const TOpClassDesc& LookupOpClass(EOpClassMethod method, ui32 typeId);
+
 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 28539412a51..91b4395791a 100644
--- a/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp
+++ b/ydb/library/yql/parser/pg_catalog/ut/catalog_ut.cpp
@@ -153,3 +153,23 @@ Y_UNIT_TEST_SUITE(TAggregationsTests) {
         UNIT_ASSERT_VALUES_EQUAL(ret.DeserializeFuncId, 0);
     }
 }
+
+Y_UNIT_TEST_SUITE(TOpClassesTests) {
+    Y_UNIT_TEST(TestMissing) {
+        UNIT_ASSERT_EXCEPTION(LookupOpClass(EOpClassMethod::Btree, LookupType("json").TypeId), yexception);
+    }
+
+   Y_UNIT_TEST(TestOk) {
+        auto ret = LookupOpClass(EOpClassMethod::Btree, LookupType("int4").TypeId);
+        UNIT_ASSERT(ret.Method == EOpClassMethod::Btree);
+        UNIT_ASSERT_VALUES_EQUAL(ret.TypeId, LookupType("int4").TypeId);
+        UNIT_ASSERT_VALUES_EQUAL(ret.Name, "int4_ops");
+        UNIT_ASSERT_VALUES_EQUAL(ret.Family, "btree/integer_ops");
+
+        ret = LookupOpClass(EOpClassMethod::Hash, LookupType("int4").TypeId);
+        UNIT_ASSERT(ret.Method == EOpClassMethod::Hash);
+        UNIT_ASSERT_VALUES_EQUAL(ret.TypeId, LookupType("int4").TypeId);
+        UNIT_ASSERT_VALUES_EQUAL(ret.Name, "int4_ops");
+        UNIT_ASSERT_VALUES_EQUAL(ret.Family, "hash/integer_ops");
+   }
+}