Skip to content

Commit 11e4668

Browse files
vitstnvitstn
authored
support of AnyNonArray type, fixes for array_agg (#1958)
1 parent 68ef320 commit 11e4668

File tree

9 files changed

+114
-14
lines changed

9 files changed

+114
-14
lines changed

‎ydb/docs/en/core/postgresql/_includes/functions.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3866,16 +3866,16 @@ Table 9.57. General-Purpose Aggregate Functions
38663866
#|
38673867
||Function|Description|Partial Mode|Example||
38683868
||array_agg ( anynonarray ) → anyarray|
3869-
Collects all the input values, including nulls, into an array. (NOT SUPPORTED)|
3869+
Collects all the input values, including nulls, into an array.|
38703870
No|
38713871
```sql
3872-
#SELECT array_agg(x) FROM (VALUES (1),(2)) a(x) → {1,2}
3872+
SELECT array_agg(x) FROM (VALUES (1),(2)) a(x) → {1,2}
38733873
```||
38743874
||array_agg ( anyarray ) → anyarray|
3875-
Concatenates all the input arrays into an array of one higher dimension. (The inputs must all have the same dimensionality, and cannot be empty or null.) (NOT SUPPORTED)|
3875+
Concatenates all the input arrays into an array of one higher dimension. (The inputs must all have the same dimensionality, and cannot be empty or null.)|
38763876
No|
38773877
```sql
3878-
#SELECT array_agg(x) FROM (VALUES (Array[1,2]),(Array[3,4])) a(x) → {{1,2},{3,4}}
3878+
SELECT array_agg(x) FROM (VALUES (Array[1,2]),(Array[3,4])) a(x) → {{1,2},{3,4}}
38793879
```||
38803880
||avg ( smallint ) → numeric
38813881
avg ( integer ) → numeric

‎ydb/docs/ru/core/postgresql/_includes/functions.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3866,16 +3866,16 @@ Table 9.57. General-Purpose Aggregate Functions
38663866
#|
38673867
||Function|Description|Partial Mode|Example||
38683868
||array_agg ( anynonarray ) → anyarray|
3869-
Collects all the input values, including nulls, into an array. (NOT SUPPORTED)|
3869+
Collects all the input values, including nulls, into an array.|
38703870
No|
38713871
```sql
3872-
#SELECT array_agg(x) FROM (VALUES (1),(2)) a(x) → {1,2}
3872+
SELECT array_agg(x) FROM (VALUES (1),(2)) a(x) → {1,2}
38733873
```||
38743874
||array_agg ( anyarray ) → anyarray|
3875-
Concatenates all the input arrays into an array of one higher dimension. (The inputs must all have the same dimensionality, and cannot be empty or null.) (NOT SUPPORTED)|
3875+
Concatenates all the input arrays into an array of one higher dimension. (The inputs must all have the same dimensionality, and cannot be empty or null.)|
38763876
No|
38773877
```sql
3878-
#SELECT array_agg(x) FROM (VALUES (Array[1,2]),(Array[3,4])) a(x) → {{1,2},{3,4}}
3878+
SELECT array_agg(x) FROM (VALUES (Array[1,2]),(Array[3,4])) a(x) → {{1,2},{3,4}}
38793879
```||
38803880
||avg ( smallint ) → numeric
38813881
avg ( integer ) → numeric

‎ydb/library/yql/core/type_ann/type_ann_pg.cpp

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ bool IsCastRequired(ui32 fromTypeId, ui32 toTypeId) {
2222
if (toTypeId == fromTypeId) {
2323
return false;
2424
}
25-
if (toTypeId == NPg::AnyOid || toTypeId == NPg::AnyArrayOid) {
25+
if (toTypeId == NPg::AnyOid || toTypeId == NPg::AnyArrayOid || toTypeId == NPg::AnyNonArrayOid) {
2626
return false;
2727
}
2828
return true;
@@ -224,6 +224,7 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
224224
}
225225

226226
bool rangeFunction = false;
227+
ui32 refinedType = 0;
227228
for (const auto& setting : input->Child(isResolved ? 2 : 1)->Children()) {
228229
if (!EnsureTupleMinSize(*setting, 1, ctx.Expr)) {
229230
return IGraphTransformer::TStatus::Error;
@@ -236,6 +237,16 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
236237
auto content = setting->Head().Content();
237238
if (content == "range") {
238239
rangeFunction = true;
240+
} else if (content == "type") {
241+
if (!EnsureTupleSize(*setting, 2, ctx.Expr)) {
242+
return IGraphTransformer::TStatus::Error;
243+
}
244+
245+
if (!EnsureAtom(setting->Tail(), ctx.Expr)) {
246+
return IGraphTransformer::TStatus::Error;
247+
}
248+
249+
refinedType = NPg::LookupType(TString(setting->Tail().Content())).TypeId;
239250
} else {
240251
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
241252
TStringBuilder() << "Unexpected setting " << content << " in function " << name));
@@ -286,9 +297,17 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
286297
return IGraphTransformer::TStatus::Error;
287298
}
288299

289-
const TTypeAnnotationNode* result = ctx.Expr.MakeType<TPgExprType>(proc.ResultType);
300+
auto resultType = proc.ResultType;
301+
AdjustReturnType(resultType, proc.ArgTypes, argTypes);
302+
if (resultType == NPg::AnyArrayOid && refinedType) {
303+
const auto& refinedDesc = NPg::LookupType(refinedType);
304+
YQL_ENSURE(refinedDesc.ArrayTypeId == refinedDesc.TypeId);
305+
resultType = refinedDesc.TypeId;
306+
}
307+
308+
const TTypeAnnotationNode* result = ctx.Expr.MakeType<TPgExprType>(resultType);
290309
TMaybe<TColumnOrder> resultColumnOrder;
291-
if (proc.ResultType == NPg::RecordOid && rangeFunction) {
310+
if (resultType == NPg::RecordOid && rangeFunction) {
292311
if (proc.OutputArgNames.empty()) {
293312
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
294313
TStringBuilder() << "Aggregate function " << name << " cannot be used in FROM"));
@@ -780,6 +799,7 @@ IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode:
780799
resultType = NPg::LookupProc(aggDesc.FinalFuncId).ResultType;
781800
}
782801

802+
AdjustReturnType(resultType, aggDesc.ArgTypes, argTypes);
783803
auto result = ctx.Expr.MakeType<TPgExprType>(resultType);
784804
input->SetTypeAnn(result);
785805

‎ydb/library/yql/core/yql_expr_type_annotation.cpp

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6347,22 +6347,42 @@ TExprNode::TPtr ExpandPgAggregationTraits(TPositionHandle pos, const NPg::TAggre
63476347
auto saveLambda = idLambda;
63486348
auto loadLambda = idLambda;
63496349
auto finishLambda = idLambda;
6350+
auto nullValue = ctx.NewCallable(pos, "Null", {});
63506351
if (aggDesc.FinalFuncId) {
6352+
const ui32 originalAggResultType = NPg::LookupProc(aggDesc.FinalFuncId).ResultType;
6353+
ui32 aggResultType = originalAggResultType;
6354+
AdjustReturnType(aggResultType, aggDesc.ArgTypes, argTypes);
63516355
finishLambda = ctx.Builder(pos)
63526356
.Lambda()
63536357
.Param("state")
63546358
.Callable("PgResolvedCallCtx")
63556359
.Atom(0, NPg::LookupProc(aggDesc.FinalFuncId).Name)
63566360
.Atom(1, ToString(aggDesc.FinalFuncId))
63576361
.List(2)
6362+
.Do([aggResultType, originalAggResultType](TExprNodeBuilder& builder) -> TExprNodeBuilder& {
6363+
if (aggResultType != originalAggResultType) {
6364+
builder.List(0)
6365+
.Atom(0, "type")
6366+
.Atom(1, NPg::LookupType(aggResultType).Name)
6367+
.Seal();
6368+
}
6369+
6370+
return builder;
6371+
})
63586372
.Seal()
63596373
.Arg(3, "state")
6374+
.Do([&aggDesc, nullValue](TExprNodeBuilder& builder) -> TExprNodeBuilder& {
6375+
if (aggDesc.FinalExtra) {
6376+
builder.Add(4, nullValue);
6377+
}
6378+
6379+
return builder;
6380+
})
63606381
.Seal()
63616382
.Seal()
63626383
.Build();
63636384
}
63646385

6365-
auto nullValue = ctx.NewCallable(pos, "Null", {});
63666386
auto initValue = nullValue;
63676387
if (aggDesc.InitValue) {
63686388
initValue = ctx.Builder(pos)
@@ -6634,6 +6654,45 @@ TExprNode::TPtr ExpandPgAggregationTraits(TPositionHandle pos, const NPg::TAggre
66346654
}
66356655
}
66366656

6657+
void AdjustReturnType(ui32& returnType, const TVector<ui32>& procArgTypes, const TVector<ui32>& argTypes) {
6658+
YQL_ENSURE(procArgTypes.size() >= argTypes.size());
6659+
if (returnType == NPg::AnyArrayOid) {
6660+
TMaybe<ui32> inputElementType;
6661+
TMaybe<ui32> inputArrayType;
6662+
for (ui32 i = 0; i < argTypes.size(); ++i) {
6663+
if (!argTypes[i]) {
6664+
continue;
6665+
}
6666+
6667+
if (procArgTypes[i] == NPg::AnyNonArrayOid) {
6668+
if (!inputElementType) {
6669+
inputElementType = argTypes[i];
6670+
} else {
6671+
if (*inputElementType != argTypes[i]) {
6672+
return;
6673+
}
6674+
}
6675+
}
6676+
6677+
if (procArgTypes[i] == NPg::AnyArrayOid) {
6678+
if (!inputArrayType) {
6679+
inputArrayType = argTypes[i];
6680+
} else {
6681+
if (*inputArrayType != argTypes[i]) {
6682+
return;
6683+
}
6684+
}
6685+
}
6686+
}
6687+
6688+
if (inputElementType) {
6689+
returnType = NPg::LookupType(*inputElementType).ArrayTypeId;
6690+
} else if (inputArrayType) {
6691+
returnType = *inputArrayType;
6692+
}
6693+
}
6694+
}
6695+
66376696
const TTypeAnnotationNode* GetOriginalResultType(TPositionHandle pos, bool isMany, const TTypeAnnotationNode* originalExtractorType, TExprContext& ctx) {
66386697
if (!EnsureStructType(pos, *originalExtractorType, ctx)) {
66396698
return nullptr;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,7 @@ bool GetMinMaxResultType(const TPositionHandle& pos, const TTypeAnnotationNode&
339339
IGraphTransformer::TStatus ExtractPgTypesFromMultiLambda(TExprNode::TPtr& lambda, TVector<ui32>& argTypes,
340340
bool& needRetype, TExprContext& ctx);
341341

342+
void AdjustReturnType(ui32& returnType, const TVector<ui32>& procArgTypes, const TVector<ui32>& argTypes);
342343
TExprNode::TPtr ExpandPgAggregationTraits(TPositionHandle pos, const NPg::TAggregateDesc& aggDesc, bool onWindow,
343344
const TExprNode::TPtr& lambda, const TVector<ui32>& argTypes, const TTypeAnnotationNode* itemType, TExprContext& ctx);
344345

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,12 @@ bool IsCompatibleTo(ui32 actualTypeId, ui32 expectedTypeId, const TTypes& types)
7373
return actualDescPtr->ArrayTypeId == actualDescPtr->TypeId;
7474
}
7575

76+
if (expectedTypeId == AnyNonArrayOid) {
77+
const auto& actualDescPtr = types.FindPtr(actualTypeId);
78+
Y_ENSURE(actualDescPtr);
79+
return actualDescPtr->ArrayTypeId != actualDescPtr->TypeId;
80+
}
81+
7682
return false;
7783
}
7884

@@ -753,6 +759,8 @@ class TAggregationsParser : public TParser {
753759
}
754760
} else if (key == "agginitval") {
755761
LastAggregation.InitValue = value;
762+
} else if (key == "aggfinalextra") {
763+
LastAggregation.FinalExtra = (value == "t");;
756764
}
757765
}
758766

@@ -1920,7 +1928,6 @@ bool IsCoercible(ui32 fromTypeId, ui32 toTypeId, ECoercionCode coercionType, con
19201928
if (toTypeId == AnyOid) {
19211929
return true;
19221930
}
1923-
19241931
//TODO: support polymorphic types
19251932

19261933
if (fromTypeId == UnknownOid) {
@@ -1943,6 +1950,12 @@ bool IsCoercible(ui32 fromTypeId, ui32 toTypeId, ECoercionCode coercionType, con
19431950
return actualDescPtr->ArrayTypeId == actualDescPtr->TypeId;
19441951
}
19451952

1953+
if (toTypeId == AnyNonArrayOid) {
1954+
const auto& actualDescPtr = catalog.Types.FindPtr(fromTypeId);
1955+
Y_ENSURE(actualDescPtr);
1956+
return actualDescPtr->ArrayTypeId != actualDescPtr->TypeId;
1957+
}
1958+
19461959
return false;
19471960
}
19481961

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ namespace NYql::NPg {
1212
constexpr ui32 UnknownOid = 705;
1313
constexpr ui32 AnyOid = 2276;
1414
constexpr ui32 AnyArrayOid = 2277;
15+
constexpr ui32 AnyNonArrayOid = 2776;
1516
constexpr ui32 RecordOid = 2249;
1617
constexpr ui32 VarcharOid = 1043;
1718
constexpr ui32 TextOid = 25;
@@ -156,6 +157,7 @@ struct TAggregateDesc {
156157
ui32 SerializeFuncId = 0;
157158
ui32 DeserializeFuncId = 0;
158159
TString InitValue;
160+
bool FinalExtra = false;
159161
};
160162

161163
enum class EAmType {

‎ydb/library/yql/parser/pg_catalog/ut/catalog_consts_ut.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ Y_UNIT_TEST_SUITE(TConstantsTests) {
5151
UNIT_ASSERT_VALUES_EQUAL(typeDesc.TypeId, VarcharOid);
5252
typeDesc = LookupType("text");
5353
UNIT_ASSERT_VALUES_EQUAL(typeDesc.TypeId, TextOid);
54+
typeDesc = LookupType("anynonarray");
55+
UNIT_ASSERT_VALUES_EQUAL(typeDesc.TypeId, AnyNonArrayOid);
5456
}
5557

5658
Y_UNIT_TEST(TRelationOidConsts) {

‎ydb/library/yql/parser/pg_wrapper/syscache.cpp

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ extern "C" {
1111
#include "catalog/pg_type_d.h"
1212
#include "catalog/pg_authid.h"
1313
#include "access/htup_details.h"
14+
#include "utils/fmgroids.h"
1415
}
1516

1617
#undef TypeName
@@ -183,7 +184,7 @@ struct TSysCache {
183184
std::fill_n(nulls, Natts_pg_type, true);
184185
std::fill_n(nulls, Anum_pg_type_typcollation, false); // fixed part of Form_pg_type
185186
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_oid, oid);
186-
auto name = MakeFixedString(desc.Name, NPg::LookupType(NAMEOID).TypeLen);
187+
auto name = MakeFixedString(desc.Name, NAMEDATALEN);
187188
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typname, (Datum)name);
188189
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typbyval, desc.PassByValue);
189190
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typlen, desc.TypeLen);
@@ -193,6 +194,8 @@ struct TSysCache {
193194
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typisdefined, true);
194195
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typdelim, desc.TypeDelim);
195196
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typarray, desc.ArrayTypeId);
197+
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typsubscript,
198+
(desc.ArrayTypeId == desc.TypeId) ? F_ARRAY_SUBSCRIPT_HANDLER : desc.TypeSubscriptFuncId);
196199
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typelem, desc.ElementTypeId);
197200
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typinput, desc.InFuncId);
198201
FillDatum(Natts_pg_type, values, nulls, Anum_pg_type_typoutput, desc.OutFuncId);

0 commit comments

Comments
 (0)