Skip to content

Commit 514df1a

Browse files
vitstnvitstn
authored
Pass pg extensions & catalog to mrjob (#7516)
1 parent 1593a2d commit 514df1a

File tree

19 files changed

+756
-36
lines changed

19 files changed

+756
-36
lines changed

‎ydb/library/yql/ast/yql_expr.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3282,6 +3282,20 @@ ui32 TPgExprType::GetFlags(ui32 typeId) {
32823282
return ret;
32833283
}
32843284

3285+
ui64 TPgExprType::GetPgExtensionsMask(ui32 typeId) {
3286+
auto descPtr = &NPg::LookupType(typeId);
3287+
return MakePgExtensionMask(descPtr->ExtensionIndex);
3288+
}
3289+
3290+
ui64 MakePgExtensionMask(ui32 extensionIndex) {
3291+
if (!extensionIndex) {
3292+
return 0;
3293+
}
3294+
3295+
YQL_ENSURE(extensionIndex <= 64);
3296+
return 1ull << (extensionIndex - 1);
3297+
}
3298+
32853299
TExprContext::TExprContext(ui64 nextUniqueId)
32863300
: StringPool(4096)
32873301
, NextUniqueId(nextUniqueId)

‎ydb/library/yql/ast/yql_expr.h

Lines changed: 47 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -152,10 +152,11 @@ void ReportError(TExprContext& ctx, const TIssue& issue);
152152
153153
class TTypeAnnotationNode {
154154
protected:
155-
TTypeAnnotationNode(ETypeAnnotationKind kind, ui32 flags, ui64 hash)
155+
TTypeAnnotationNode(ETypeAnnotationKind kind, ui32 flags, ui64 hash, ui64 usedPgExtensions)
156156
: Kind(kind)
157157
, Flags(flags)
158158
, Hash(hash)
159+
, UsedPgExtensions(usedPgExtensions)
159160
{
160161
}
161162
@@ -278,6 +279,10 @@ class TTypeAnnotationNode {
278279
return Hash;
279280
}
280281
282+
ui64 GetUsedPgExtensions() const {
283+
return UsedPgExtensions;
284+
}
285+
281286
bool Equals(const TTypeAnnotationNode& node) const;
282287
void Accept(TTypeAnnotationVisitor& visitor) const;
283288
@@ -310,10 +315,21 @@ class TTypeAnnotationNode {
310315
return flags;
311316
}
312317
318+
template <typename T>
319+
static ui64 CombinePgExtensions(const T& items) {
320+
ui64 mask = 0;
321+
for (auto& item : items) {
322+
mask |= item->GetUsedPgExtensions();
323+
}
324+
325+
return mask;
326+
}
327+
313328
private:
314329
const ETypeAnnotationKind Kind;
315330
const ui32 Flags;
316331
const ui64 Hash;
332+
const ui64 UsedPgExtensions;
317333
};
318334
319335
class TUnitExprType : public TTypeAnnotationNode {
@@ -322,7 +338,7 @@ class TUnitExprType : public TTypeAnnotationNode {
322338
323339
TUnitExprType(ui64 hash)
324340
: TTypeAnnotationNode(KindValue,
325-
TypeNonComputable | TypeNonPersistable, hash)
341+
TypeNonComputable | TypeNonPersistable, hash, 0)
326342
{
327343
}
328344
@@ -341,7 +357,7 @@ class TTupleExprType : public TTypeAnnotationNode {
341357
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Tuple;
342358
343359
TTupleExprType(ui64 hash, const TTypeAnnotationNode::TListType& items)
344-
: TTypeAnnotationNode(KindValue, CombineFlags(items), hash)
360+
: TTypeAnnotationNode(KindValue, CombineFlags(items), hash, CombinePgExtensions(items))
345361
, Items(items)
346362
{
347363
}
@@ -390,7 +406,7 @@ class TMultiExprType : public TTypeAnnotationNode {
390406
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Multi;
391407
392408
TMultiExprType(ui64 hash, const TTypeAnnotationNode::TListType& items)
393-
: TTypeAnnotationNode(KindValue, CombineFlags(items), hash)
409+
: TTypeAnnotationNode(KindValue, CombineFlags(items), hash, CombinePgExtensions(items))
394410
, Items(items)
395411
{
396412
}
@@ -445,7 +461,7 @@ class TItemExprType : public TTypeAnnotationNode {
445461
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Item;
446462
447463
TItemExprType(ui64 hash, const TStringBuf& name, const TTypeAnnotationNode* itemType)
448-
: TTypeAnnotationNode(KindValue, itemType->GetFlags(), hash)
464+
: TTypeAnnotationNode(KindValue, itemType->GetFlags(), hash, itemType->GetUsedPgExtensions())
449465
, Name(name)
450466
, ItemType(itemType)
451467
{
@@ -501,7 +517,7 @@ class TStructExprType : public TTypeAnnotationNode {
501517
};
502518
503519
TStructExprType(ui64 hash, const TVector<const TItemExprType*>& items)
504-
: TTypeAnnotationNode(KindValue, TypeNonComparable | CombineFlags(items), hash)
520+
: TTypeAnnotationNode(KindValue, TypeNonComparable | CombineFlags(items), hash, CombinePgExtensions(items))
505521
, Items(items)
506522
{
507523
}
@@ -623,7 +639,7 @@ class TListExprType : public TTypeAnnotationNode {
623639
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::List;
624640
625641
TListExprType(ui64 hash, const TTypeAnnotationNode* itemType)
626-
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeHasDynamicSize, hash)
642+
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeHasDynamicSize, hash, itemType->GetUsedPgExtensions())
627643
, ItemType(itemType)
628644
{
629645
}
@@ -650,7 +666,7 @@ class TStreamExprType : public TTypeAnnotationNode {
650666
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Stream;
651667
652668
TStreamExprType(ui64 hash, const TTypeAnnotationNode* itemType)
653-
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
669+
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash, itemType->GetUsedPgExtensions())
654670
, ItemType(itemType)
655671
{
656672
}
@@ -677,7 +693,7 @@ class TFlowExprType : public TTypeAnnotationNode {
677693
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Flow;
678694
679695
TFlowExprType(ui64 hash, const TTypeAnnotationNode* itemType)
680-
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
696+
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash, itemType->GetUsedPgExtensions())
681697
, ItemType(itemType)
682698
{
683699
}
@@ -704,7 +720,7 @@ class TBlockExprType : public TTypeAnnotationNode {
704720
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Block;
705721
706722
TBlockExprType(ui64 hash, const TTypeAnnotationNode* itemType)
707-
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
723+
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash, itemType->GetUsedPgExtensions())
708724
, ItemType(itemType)
709725
{
710726
}
@@ -731,7 +747,7 @@ class TScalarExprType : public TTypeAnnotationNode {
731747
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Scalar;
732748
733749
TScalarExprType(ui64 hash, const TTypeAnnotationNode* itemType)
734-
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
750+
: TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash, itemType->GetUsedPgExtensions())
735751
, ItemType(itemType)
736752
{
737753
}
@@ -758,7 +774,7 @@ class TDataExprType : public TTypeAnnotationNode {
758774
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Data;
759775
760776
TDataExprType(ui64 hash, EDataSlot slot)
761-
: TTypeAnnotationNode(KindValue, GetFlags(slot), hash)
777+
: TTypeAnnotationNode(KindValue, GetFlags(slot), hash, 0)
762778
, Slot(slot)
763779
{
764780
}
@@ -853,7 +869,7 @@ class TPgExprType : public TTypeAnnotationNode {
853869
854870
// TODO: TypeHasDynamicSize for Pg types
855871
TPgExprType(ui64 hash, ui32 typeId)
856-
: TTypeAnnotationNode(KindValue, GetFlags(typeId), hash)
872+
: TTypeAnnotationNode(KindValue, GetFlags(typeId), hash, GetPgExtensionsMask(typeId))
857873
, TypeId(typeId)
858874
{
859875
}
@@ -875,18 +891,21 @@ class TPgExprType : public TTypeAnnotationNode {
875891
876892
private:
877893
ui32 GetFlags(ui32 typeId);
894+
ui64 GetPgExtensionsMask(ui32 typeId);
878895
879896
private:
880897
ui32 TypeId;
881898
};
882899
900+
ui64 MakePgExtensionMask(ui32 extensionIndex);
901+
883902
class TWorldExprType : public TTypeAnnotationNode {
884903
public:
885904
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::World;
886905
887906
TWorldExprType(ui64 hash)
888907
: TTypeAnnotationNode(KindValue,
889-
TypeNonComposable | TypeNonComputable | TypeNonPersistable | TypeNonInspectable, hash)
908+
TypeNonComposable | TypeNonComputable | TypeNonPersistable | TypeNonInspectable, hash, 0)
890909
{
891910
}
892911
@@ -905,7 +924,7 @@ class TOptionalExprType : public TTypeAnnotationNode {
905924
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Optional;
906925
907926
TOptionalExprType(ui64 hash, const TTypeAnnotationNode* itemType)
908-
: TTypeAnnotationNode(KindValue, GetFlags(itemType), hash)
927+
: TTypeAnnotationNode(KindValue, GetFlags(itemType), hash, itemType->GetUsedPgExtensions())
909928
, ItemType(itemType)
910929
{
911930
}
@@ -945,7 +964,7 @@ class TVariantExprType : public TTypeAnnotationNode {
945964
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Variant;
946965
947966
TVariantExprType(ui64 hash, const TTypeAnnotationNode* underlyingType)
948-
: TTypeAnnotationNode(KindValue, MakeFlags(underlyingType), hash)
967+
: TTypeAnnotationNode(KindValue, MakeFlags(underlyingType), hash, underlyingType->GetUsedPgExtensions())
949968
, UnderlyingType(underlyingType)
950969
{
951970
}
@@ -977,7 +996,7 @@ class TTypeExprType : public TTypeAnnotationNode {
977996
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Type;
978997
979998
TTypeExprType(ui64 hash, const TTypeAnnotationNode* type)
980-
: TTypeAnnotationNode(KindValue, TypeNonPersistable | TypeNonComputable, hash)
999+
: TTypeAnnotationNode(KindValue, TypeNonPersistable | TypeNonComputable, hash, 0)
9811000
, Type(type)
9821001
{
9831002
}
@@ -1005,7 +1024,8 @@ class TDictExprType : public TTypeAnnotationNode {
10051024
10061025
TDictExprType(ui64 hash, const TTypeAnnotationNode* keyType, const TTypeAnnotationNode* payloadType)
10071026
: TTypeAnnotationNode(KindValue, TypeNonComparable | TypeHasDynamicSize |
1008-
keyType->GetFlags() | payloadType->GetFlags(), hash)
1027+
keyType->GetFlags() | payloadType->GetFlags(), hash,
1028+
keyType->GetUsedPgExtensions() | payloadType->GetUsedPgExtensions())
10091029
, KeyType(keyType)
10101030
, PayloadType(payloadType)
10111031
{
@@ -1042,7 +1062,7 @@ class TVoidExprType : public TTypeAnnotationNode {
10421062
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Void;
10431063
10441064
TVoidExprType(ui64 hash)
1045-
: TTypeAnnotationNode(KindValue, 0, hash)
1065+
: TTypeAnnotationNode(KindValue, 0, hash, 0)
10461066
{
10471067
}
10481068
@@ -1061,7 +1081,7 @@ class TNullExprType : public TTypeAnnotationNode {
10611081
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Null;
10621082
10631083
TNullExprType(ui64 hash)
1064-
: TTypeAnnotationNode(KindValue, TypeHasNull, hash)
1084+
: TTypeAnnotationNode(KindValue, TypeHasNull, hash, 0)
10651085
{
10661086
}
10671087
@@ -1101,7 +1121,7 @@ class TCallableExprType : public TTypeAnnotationNode {
11011121
11021122
TCallableExprType(ui64 hash, const TTypeAnnotationNode* returnType, const TVector<TArgumentInfo>& arguments
11031123
, size_t optionalArgumentsCount, const TStringBuf& payload)
1104-
: TTypeAnnotationNode(KindValue, MakeFlags(returnType), hash)
1124+
: TTypeAnnotationNode(KindValue, MakeFlags(returnType), hash, returnType->GetUsedPgExtensions())
11051125
, ReturnType(returnType)
11061126
, Arguments(arguments)
11071127
, OptionalArgumentsCount(optionalArgumentsCount)
@@ -1207,7 +1227,7 @@ class TGenericExprType : public TTypeAnnotationNode {
12071227
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Generic;
12081228
12091229
TGenericExprType(ui64 hash)
1210-
: TTypeAnnotationNode(KindValue, TypeNonComputable, hash)
1230+
: TTypeAnnotationNode(KindValue, TypeNonComputable, hash, 0)
12111231
{
12121232
}
12131233
@@ -1226,7 +1246,7 @@ class TResourceExprType : public TTypeAnnotationNode {
12261246
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Resource;
12271247
12281248
TResourceExprType(ui64 hash, const TStringBuf& tag)
1229-
: TTypeAnnotationNode(KindValue, TypeNonPersistable | TypeHasManyValues, hash)
1249+
: TTypeAnnotationNode(KindValue, TypeNonPersistable | TypeHasManyValues, hash, 0)
12301250
, Tag(tag)
12311251
{}
12321252
@@ -1253,7 +1273,7 @@ class TTaggedExprType : public TTypeAnnotationNode {
12531273
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Tagged;
12541274
12551275
TTaggedExprType(ui64 hash, const TTypeAnnotationNode* baseType, const TStringBuf& tag)
1256-
: TTypeAnnotationNode(KindValue, baseType->GetFlags(), hash)
1276+
: TTypeAnnotationNode(KindValue, baseType->GetFlags(), hash, baseType->GetUsedPgExtensions())
12571277
, BaseType(baseType)
12581278
, Tag(tag)
12591279
{}
@@ -1290,7 +1310,7 @@ class TErrorExprType : public TTypeAnnotationNode {
12901310
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Error;
12911311
12921312
TErrorExprType(ui64 hash, const TIssue& error)
1293-
: TTypeAnnotationNode(KindValue, 0, hash)
1313+
: TTypeAnnotationNode(KindValue, 0, hash, 0)
12941314
, Error(error)
12951315
{}
12961316
@@ -1315,7 +1335,7 @@ class TEmptyListExprType : public TTypeAnnotationNode {
13151335
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::EmptyList;
13161336
13171337
TEmptyListExprType(ui64 hash)
1318-
: TTypeAnnotationNode(KindValue, 0, hash)
1338+
: TTypeAnnotationNode(KindValue, 0, hash, 0)
13191339
{
13201340
}
13211341
@@ -1334,7 +1354,7 @@ class TEmptyDictExprType : public TTypeAnnotationNode {
13341354
static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::EmptyDict;
13351355
13361356
TEmptyDictExprType(ui64 hash)
1337-
: TTypeAnnotationNode(KindValue, 0, hash)
1357+
: TTypeAnnotationNode(KindValue, 0, hash, 0)
13381358
{
13391359
}
13401360

‎ydb/library/yql/core/pg_ext/yql_pg_ext.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ void PgExtensionsFromProto(const NYql::NProto::TPgExtensions& proto,
1515

1616
desc.LibraryPath = e.GetLibraryPath();
1717
desc.TypesOnly = e.GetTypesOnly();
18+
desc.LibraryMD5 = e.GetLibraryMD5();
1819
extensions.emplace_back(desc);
1920
}
2021
}

‎ydb/library/yql/core/yql_user_data.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ enum class EUserDataBlockUsage {
2525
Content,
2626
Udf,
2727
Library,
28+
PgExt,
2829
End,
2930
};
3031
typedef TEnumBitSet<EUserDataBlockUsage, static_cast<int>(EUserDataBlockUsage::Begin),

0 commit comments

Comments
 (0)