Skip to content

Commit bbf4a41

Browse files
vitstnvitstn
authored
Support of OPERATOR CLASS for PG extensions (#8031)
1 parent 0f0f3bc commit bbf4a41

File tree

5 files changed

+486
-36
lines changed

5 files changed

+486
-36
lines changed

‎ydb/library/yql/parser/pg_catalog/catalog.cpp

Lines changed: 213 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,7 @@ using TAms = THashMap<ui32, TAmDesc>;
5959

6060
using TNamespaces = THashMap<decltype(TNamespaceDesc::Oid), TNamespaceDesc>;
6161

62-
// We parse OpFamilies' IDs for now. If we ever needed oid_symbol,
63-
// create TOpFamilyDesc class alike other catalogs
64-
using TOpFamilies = THashMap<TString, ui32>;
62+
using TOpFamilies = THashMap<TString, TOpFamilyDesc>;
6563

6664
using TOpClasses = THashMap<std::pair<EOpClassMethod, ui32>, TOpClassDesc>;
6765

@@ -1073,10 +1071,11 @@ class TOpFamiliesParser : public TParser {
10731071
if (IsSupported) {
10741072
Y_ENSURE(LastOpfId != InvalidOid);
10751073

1076-
// TODO: log or throw if dict keys aren't unique
10771074
// opfamily references have opf_method/opf_name format in PG catalogs
1078-
OpFamilies[LastOpfMethod.append('/').append(
1079-
LastOpfName)] = LastOpfId; // LastOpfMethod is modified here. Use with caution till its reinit
1075+
TOpFamilyDesc desc;
1076+
desc.Name = LastOpfMethod + "/" + LastOpfName;
1077+
desc.FamilyId = LastOpfId;
1078+
Y_ENSURE(OpFamilies.emplace(desc.Name, desc).second);
10801079
}
10811080

10821081
IsSupported = true;
@@ -1124,7 +1123,7 @@ class TOpClassesParser : public TParser {
11241123
auto opFamilyPtr = OpFamilies.FindPtr(value);
11251124

11261125
if (opFamilyPtr) {
1127-
LastOpClass.FamilyId = *opFamilyPtr;
1126+
LastOpClass.FamilyId = opFamilyPtr->FamilyId;
11281127
} else {
11291128
IsSupported = false;
11301129
}
@@ -1181,7 +1180,7 @@ class TAmOpsParser : public TParser {
11811180
LastAmOp.Family = value;
11821181
auto opFamilyPtr = OpFamilies.FindPtr(value);
11831182
if (opFamilyPtr) {
1184-
LastAmOp.FamilyId = *opFamilyPtr;
1183+
LastAmOp.FamilyId = opFamilyPtr->FamilyId;
11851184
} else {
11861185
IsSupported = false;
11871186
}
@@ -1292,7 +1291,7 @@ class TAmProcsParser : public TParser {
12921291
auto opFamilyPtr = OpFamilies.FindPtr(value);
12931292

12941293
if (opFamilyPtr) {
1295-
LastAmProc.FamilyId = *opFamilyPtr;
1294+
LastAmProc.FamilyId = opFamilyPtr->FamilyId;
12961295
} else {
12971296
IsSupported = false;
12981297
}
@@ -1843,10 +1842,10 @@ struct TCatalog : public IExtensionSqlBuilder {
18431842
State->AggregationsByName[v.Name].push_back(k);
18441843
}
18451844

1846-
TOpFamilies opFamilies = ParseOpFamilies(opFamiliesData);
1847-
State->OpClasses = ParseOpClasses(opClassData, State->TypeByName, opFamilies);
1848-
State->AmOps = ParseAmOps(amOpData, State->TypeByName, State->Types, State->OperatorsByName, State->Operators, opFamilies);
1849-
State->AmProcs = ParseAmProcs(amProcData, State->TypeByName, State->ProcByName, State->Procs, opFamilies);
1845+
State->OpFamilies = ParseOpFamilies(opFamiliesData);
1846+
State->OpClasses = ParseOpClasses(opClassData, State->TypeByName, State->OpFamilies);
1847+
State->AmOps = ParseAmOps(amOpData, State->TypeByName, State->Types, State->OperatorsByName, State->Operators, State->OpFamilies);
1848+
State->AmProcs = ParseAmProcs(amProcData, State->TypeByName, State->ProcByName, State->Procs, State->OpFamilies);
18501849
State->Ams = ParseAms(amData);
18511850
State->Namespaces = FillNamespaces();
18521851
for (auto& [k, v] : State->Types) {
@@ -1855,30 +1854,8 @@ struct TCatalog : public IExtensionSqlBuilder {
18551854
if (regClasses.contains(lookupId)) {
18561855
lookupId = OidOid;
18571856
}
1858-
auto btreeOpClassPtr = State->OpClasses.FindPtr(std::make_pair(EOpClassMethod::Btree, lookupId));
1859-
if (btreeOpClassPtr) {
1860-
auto lessAmOpPtr = State->AmOps.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmStrategy::Less), lookupId, lookupId));
1861-
Y_ENSURE(lessAmOpPtr);
1862-
auto equalAmOpPtr = State->AmOps.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmStrategy::Equal), lookupId, lookupId));
1863-
Y_ENSURE(equalAmOpPtr);
1864-
auto lessOperPtr = State->Operators.FindPtr(lessAmOpPtr->OperId);
1865-
Y_ENSURE(lessOperPtr);
1866-
auto equalOperPtr = State->Operators.FindPtr(equalAmOpPtr->OperId);
1867-
Y_ENSURE(equalOperPtr);
1868-
v.LessProcId = lessOperPtr->ProcId;
1869-
v.EqualProcId = equalOperPtr->ProcId;
1870-
1871-
auto compareAmProcPtr = State->AmProcs.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmProcNum::Compare), lookupId, lookupId));
1872-
Y_ENSURE(compareAmProcPtr);
1873-
v.CompareProcId = compareAmProcPtr->ProcId;
1874-
}
18751857

1876-
auto hashOpClassPtr = State->OpClasses.FindPtr(std::make_pair(EOpClassMethod::Hash, lookupId));
1877-
if (hashOpClassPtr) {
1878-
auto hashAmProcPtr = State->AmProcs.FindPtr(std::make_tuple(hashOpClassPtr->FamilyId, ui32(EHashAmProcNum::Hash), lookupId, lookupId));
1879-
Y_ENSURE(hashAmProcPtr);
1880-
v.HashProcId = hashAmProcPtr->ProcId;
1881-
}
1858+
CacheAmFuncs(lookupId, v);
18821859
}
18831860
}
18841861

@@ -1956,6 +1933,33 @@ struct TCatalog : public IExtensionSqlBuilder {
19561933
}
19571934
}
19581935

1936+
void CacheAmFuncs(ui32 typeId, TTypeDesc& v) {
1937+
auto btreeOpClassPtr = State->OpClasses.FindPtr(std::make_pair(EOpClassMethod::Btree, typeId));
1938+
if (btreeOpClassPtr) {
1939+
auto lessAmOpPtr = State->AmOps.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmStrategy::Less), typeId, typeId));
1940+
Y_ENSURE(lessAmOpPtr);
1941+
auto equalAmOpPtr = State->AmOps.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmStrategy::Equal), typeId, typeId));
1942+
Y_ENSURE(equalAmOpPtr);
1943+
auto lessOperPtr = State->Operators.FindPtr(lessAmOpPtr->OperId);
1944+
Y_ENSURE(lessOperPtr);
1945+
auto equalOperPtr = State->Operators.FindPtr(equalAmOpPtr->OperId);
1946+
Y_ENSURE(equalOperPtr);
1947+
v.LessProcId = lessOperPtr->ProcId;
1948+
v.EqualProcId = equalOperPtr->ProcId;
1949+
1950+
auto compareAmProcPtr = State->AmProcs.FindPtr(std::make_tuple(btreeOpClassPtr->FamilyId, ui32(EBtreeAmProcNum::Compare), typeId, typeId));
1951+
Y_ENSURE(compareAmProcPtr);
1952+
v.CompareProcId = compareAmProcPtr->ProcId;
1953+
}
1954+
1955+
auto hashOpClassPtr = State->OpClasses.FindPtr(std::make_pair(EOpClassMethod::Hash, typeId));
1956+
if (hashOpClassPtr) {
1957+
auto hashAmProcPtr = State->AmProcs.FindPtr(std::make_tuple(hashOpClassPtr->FamilyId, ui32(EHashAmProcNum::Hash), typeId, typeId));
1958+
Y_ENSURE(hashAmProcPtr);
1959+
v.HashProcId = hashAmProcPtr->ProcId;
1960+
}
1961+
}
1962+
19591963
void CreateProc(const TProcDesc& desc) final {
19601964
Y_ENSURE(desc.ExtensionIndex);
19611965
TProcDesc newDesc = desc;
@@ -2183,6 +2187,52 @@ struct TCatalog : public IExtensionSqlBuilder {
21832187
State->AllowedProcs.insert(State->Procs.FindPtr(desc.TransFuncId)->Name);
21842188
}
21852189

2190+
void CreateOpClass(const TOpClassDesc& opclass, const TVector<TAmOpDesc>& ops, const TVector<TAmProcDesc>& procs) final {
2191+
Y_ENSURE(opclass.ExtensionIndex);
2192+
auto newDesc = opclass;
2193+
newDesc.Family = to_lower(newDesc.Family);
2194+
newDesc.Name = to_lower(newDesc.Name);
2195+
auto newFamilyId = 16000 + State->OpFamilies.size();
2196+
TOpFamilyDesc opFamilyDesc;
2197+
opFamilyDesc.FamilyId = newFamilyId;
2198+
opFamilyDesc.Name = newDesc.Family;
2199+
opFamilyDesc.ExtensionIndex = opclass.ExtensionIndex;
2200+
Y_ENSURE(State->OpFamilies.emplace(newDesc.Family, opFamilyDesc).second);
2201+
newDesc.FamilyId = newFamilyId;
2202+
const auto key = std::make_pair(newDesc.Method, newDesc.TypeId);
2203+
Y_ENSURE(State->OpClasses.emplace(key, newDesc).second);
2204+
for (const auto& o : ops) {
2205+
Y_ENSURE(opclass.ExtensionIndex == o.ExtensionIndex);
2206+
Y_ENSURE(opclass.Family == o.Family);
2207+
auto newOpDesc = o;
2208+
newOpDesc.FamilyId = newFamilyId;
2209+
newOpDesc.Family = newDesc.Name;
2210+
Y_ENSURE(State->AmOps.emplace(std::make_tuple(newFamilyId, o.Strategy, o.LeftType, o.RightType), newOpDesc).second);
2211+
auto operPtr = State->Operators.FindPtr(o.OperId);
2212+
Y_ENSURE(operPtr);
2213+
auto procPtr = State->Procs.FindPtr(operPtr->ProcId);
2214+
Y_ENSURE(procPtr);
2215+
State->AllowedProcs.emplace(procPtr->Name);
2216+
}
2217+
2218+
for (const auto& p : procs) {
2219+
Y_ENSURE(opclass.ExtensionIndex == p.ExtensionIndex);
2220+
Y_ENSURE(opclass.Family == p.Family);
2221+
auto newProcDesc = p;
2222+
newProcDesc.FamilyId = newFamilyId;
2223+
newProcDesc.Family = newDesc.Name;
2224+
Y_ENSURE(State->AmProcs.emplace(std::make_tuple(newFamilyId, p.ProcNum, p.LeftType, p.RightType), newProcDesc).second);
2225+
auto procPtr = State->Procs.FindPtr(p.ProcId);
2226+
Y_ENSURE(procPtr);
2227+
State->AllowedProcs.emplace(procPtr->Name);
2228+
}
2229+
2230+
auto typePtr = State->Types.FindPtr(opclass.TypeId);
2231+
Y_ENSURE(typePtr);
2232+
Y_ENSURE(typePtr->ExtensionIndex == opclass.ExtensionIndex);
2233+
CacheAmFuncs(opclass.TypeId, *typePtr);
2234+
}
2235+
21862236
static const TCatalog& Instance() {
21872237
return *Singleton<TCatalog>();
21882238
}
@@ -2202,6 +2252,7 @@ struct TCatalog : public IExtensionSqlBuilder {
22022252
TAggregations Aggregations;
22032253
TAms Ams;
22042254
TNamespaces Namespaces;
2255+
TOpFamilies OpFamilies;
22052256
TOpClasses OpClasses;
22062257
TAmOps AmOps;
22072258
TAmProcs AmProcs;
@@ -3840,6 +3891,87 @@ TString ExportExtensions(const TMaybe<TSet<ui32>>& filter) {
38403891
protoAggregation->SetNumDirectArgs(desc.NumDirectArgs);
38413892
}
38423893

3894+
TVector<TString> extOpFamilies;
3895+
for (const auto& f : catalog.State->OpFamilies) {
3896+
const auto& desc = f.second;
3897+
if (!desc.ExtensionIndex) {
3898+
continue;
3899+
}
3900+
3901+
extOpFamilies.push_back(f.first);
3902+
}
3903+
3904+
Sort(extOpFamilies);
3905+
for (const auto& f : extOpFamilies) {
3906+
const auto& desc = *catalog.State->OpFamilies.FindPtr(f);
3907+
auto protoOpClassFamily = proto.AddOpClassFamily();
3908+
protoOpClassFamily->SetFamilyId(desc.FamilyId);
3909+
protoOpClassFamily->SetName(desc.Name);
3910+
protoOpClassFamily->SetExtensionIndex(desc.ExtensionIndex);
3911+
}
3912+
3913+
TVector<std::pair<NPg::EOpClassMethod, ui32>> extOpClasses;
3914+
for (const auto& c : catalog.State->OpClasses) {
3915+
const auto& desc = c.second;
3916+
if (!desc.ExtensionIndex) {
3917+
continue;
3918+
}
3919+
3920+
extOpClasses.push_back(c.first);
3921+
}
3922+
3923+
for (const auto& c : extOpClasses) {
3924+
const auto& desc = *catalog.State->OpClasses.FindPtr(c);
3925+
auto protoOpClass = proto.AddOpClass();
3926+
protoOpClass->SetMethod((ui32)desc.Method);
3927+
protoOpClass->SetTypeId(desc.TypeId);
3928+
protoOpClass->SetExtensionIndex(desc.ExtensionIndex);
3929+
protoOpClass->SetName(desc.Name);
3930+
protoOpClass->SetFamilyId(desc.FamilyId);
3931+
}
3932+
3933+
TVector<std::tuple<ui32, ui32, ui32, ui32>> extAmOps;
3934+
for (const auto& o : catalog.State->AmOps) {
3935+
const auto& desc = o.second;
3936+
if (!desc.ExtensionIndex) {
3937+
continue;
3938+
}
3939+
3940+
extAmOps.push_back(o.first);
3941+
}
3942+
3943+
for (const auto& o : extAmOps) {
3944+
const auto& desc = *catalog.State->AmOps.FindPtr(o);
3945+
auto protoAmOp = proto.AddAmOp();
3946+
protoAmOp->SetFamilyId(desc.FamilyId);
3947+
protoAmOp->SetStrategy(desc.Strategy);
3948+
protoAmOp->SetLeftType(desc.LeftType);
3949+
protoAmOp->SetRightType(desc.RightType);
3950+
protoAmOp->SetOperId(desc.OperId);
3951+
protoAmOp->SetExtensionIndex(desc.ExtensionIndex);
3952+
}
3953+
3954+
TVector<std::tuple<ui32, ui32, ui32, ui32>> extAmProcs;
3955+
for (const auto& p : catalog.State->AmProcs) {
3956+
const auto& desc = p.second;
3957+
if (!desc.ExtensionIndex) {
3958+
continue;
3959+
}
3960+
3961+
extAmProcs.push_back(p.first);
3962+
}
3963+
3964+
for (const auto& p : extAmProcs) {
3965+
const auto& desc = *catalog.State->AmProcs.FindPtr(p);
3966+
auto protoAmProc = proto.AddAmProc();
3967+
protoAmProc->SetFamilyId(desc.FamilyId);
3968+
protoAmProc->SetProcNum(desc.ProcNum);
3969+
protoAmProc->SetLeftType(desc.LeftType);
3970+
protoAmProc->SetRightType(desc.RightType);
3971+
protoAmProc->SetProcId(desc.ProcId);
3972+
protoAmProc->SetExtensionIndex(desc.ExtensionIndex);
3973+
}
3974+
38433975
return proto.SerializeAsString();
38443976
}
38453977

@@ -4017,6 +4149,51 @@ void ImportExtensions(const TString& exported, bool typesOnly, IExtensionLoader*
40174149
catalog.State->AggregationsByName[desc.Name].push_back(desc.AggId);
40184150
}
40194151

4152+
THashMap<ui32, TString> opFamiliesByOid;
4153+
for (const auto& protoOpClassFamily : proto.GetOpClassFamily()) {
4154+
TOpFamilyDesc desc;
4155+
desc.FamilyId = protoOpClassFamily.GetFamilyId();
4156+
desc.Name = protoOpClassFamily.GetName();
4157+
desc.ExtensionIndex = protoOpClassFamily.GetExtensionIndex();
4158+
Y_ENSURE(catalog.State->OpFamilies.emplace(desc.Name, desc).second);
4159+
Y_ENSURE(opFamiliesByOid.emplace(desc.FamilyId, desc.Name).second);
4160+
}
4161+
4162+
for (const auto& protoOpClass : proto.GetOpClass()) {
4163+
TOpClassDesc desc;
4164+
desc.Method = (EOpClassMethod)protoOpClass.GetMethod();
4165+
desc.TypeId = protoOpClass.GetTypeId();
4166+
desc.ExtensionIndex = protoOpClass.GetExtensionIndex();
4167+
desc.FamilyId = protoOpClass.GetFamilyId();
4168+
desc.Name = protoOpClass.GetName();
4169+
desc.Family = *opFamiliesByOid.FindPtr(desc.FamilyId);
4170+
Y_ENSURE(catalog.State->OpClasses.emplace(std::make_pair(desc.Method, desc.TypeId), desc).second);
4171+
}
4172+
4173+
for (const auto& protoAmOp : proto.GetAmOp()) {
4174+
TAmOpDesc desc;
4175+
desc.FamilyId = protoAmOp.GetFamilyId();
4176+
desc.Strategy = protoAmOp.GetStrategy();
4177+
desc.LeftType = protoAmOp.GetLeftType();
4178+
desc.RightType = protoAmOp.GetRightType();
4179+
desc.OperId = protoAmOp.GetOperId();
4180+
desc.ExtensionIndex = protoAmOp.GetExtensionIndex();
4181+
desc.Family = *opFamiliesByOid.FindPtr(desc.FamilyId);
4182+
Y_ENSURE(catalog.State->AmOps.emplace(std::make_tuple(desc.FamilyId, desc.Strategy, desc.LeftType, desc.RightType), desc).second);
4183+
}
4184+
4185+
for (const auto& protoAmProc : proto.GetAmProc()) {
4186+
TAmProcDesc desc;
4187+
desc.FamilyId = protoAmProc.GetFamilyId();
4188+
desc.ProcNum = protoAmProc.GetProcNum();
4189+
desc.LeftType = protoAmProc.GetLeftType();
4190+
desc.RightType = protoAmProc.GetRightType();
4191+
desc.ProcId = protoAmProc.GetProcId();
4192+
desc.ExtensionIndex = protoAmProc.GetExtensionIndex();
4193+
desc.Family = *opFamiliesByOid.FindPtr(desc.FamilyId);
4194+
Y_ENSURE(catalog.State->AmProcs.emplace(std::make_tuple(desc.FamilyId, desc.ProcNum, desc.LeftType, desc.RightType), desc).second);
4195+
}
4196+
40204197
if (!typesOnly && loader) {
40214198
for (ui32 extensionIndex = 1; extensionIndex <= catalog.State->Extensions.size(); ++extensionIndex) {
40224199
const auto& e = catalog.State->Extensions[extensionIndex - 1];

‎ydb/library/yql/parser/pg_catalog/catalog.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,12 +204,19 @@ enum class EOpClassMethod {
204204
Hash
205205
};
206206

207+
struct TOpFamilyDesc {
208+
TString Name;
209+
ui32 FamilyId = 0;
210+
ui32 ExtensionIndex = 0;
211+
};
212+
207213
struct TOpClassDesc {
208214
EOpClassMethod Method = EOpClassMethod::Btree;
209215
ui32 TypeId = 0;
210216
TString Name;
211217
TString Family;
212218
ui32 FamilyId = 0;
219+
ui32 ExtensionIndex = 0;
213220
};
214221

215222
struct TAmOpDesc {
@@ -219,6 +226,7 @@ struct TAmOpDesc {
219226
ui32 LeftType = 0;
220227
ui32 RightType = 0;
221228
ui32 OperId = 0;
229+
ui32 ExtensionIndex = 0;
222230
};
223231

224232
enum class EBtreeAmStrategy {
@@ -236,6 +244,7 @@ struct TAmProcDesc {
236244
ui32 LeftType = 0;
237245
ui32 RightType = 0;
238246
ui32 ProcId = 0;
247+
ui32 ExtensionIndex = 0;
239248
};
240249

241250
enum class EBtreeAmProcNum {
@@ -410,6 +419,8 @@ class IExtensionSqlBuilder {
410419
virtual void UpdateOper(const TOperDesc& desc) = 0;
411420

412421
virtual void CreateAggregate(const TAggregateDesc& desc) = 0;
422+
423+
virtual void CreateOpClass(const TOpClassDesc& opclass, const TVector<TAmOpDesc>& ops, const TVector<TAmProcDesc>& procs) = 0;
413424
};
414425

415426
class IExtensionSqlParser {

0 commit comments

Comments
 (0)