Skip to content

Commit 300e375

Browse files
vitstnvitstn
authored
Support of var array args (#2828)
1 parent aed1c30 commit 300e375

File tree

11 files changed

+152
-10
lines changed

11 files changed

+152
-10
lines changed

‎ydb/library/yql/parser/pg_wrapper/arrow.h‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ struct TPgKernelState : arrow::compute::KernelState {
3333
bool IsFixedResult;
3434
i32 TypeLen;
3535
std::shared_ptr<void> FmgrDataHolder;
36+
const NPg::TProcDesc* ProcDesc;
3637
};
3738

3839
template <PGFunction PgFunc>
@@ -356,7 +357,7 @@ struct TGenericExec {
356357
fcinfo->context = state.context;
357358
fcinfo->resultinfo = state.resultinfo;
358359
fcinfo->fncollation = state.fncollation;
359-
fcinfo->nargs = batch.values.size();
360+
fcinfo->nargs = batch.values.size();
360361

361362
TInputArgsAccessor<TArgsPolicy> inputArgsAccessor;
362363
inputArgsAccessor.Bind(batch.values);
@@ -371,6 +372,7 @@ struct TGenericExec {
371372

372373
for (size_t i = 0; i < length; ++i) {
373374
Datum ret;
375+
bool needToFree = false;
374376
if constexpr (!TArgsPolicy::VarArgs) {
375377
if (!constexpr_for_tuple([&](auto const& j, auto const& v) {
376378
NullableDatum d;
@@ -438,7 +440,17 @@ struct TGenericExec {
438440
}
439441

440442
fcinfo->isnull = false;
443+
if constexpr (TArgsPolicy::VarArgs) {
444+
needToFree = PrepareVariadicArray(*fcinfo, *state.ProcDesc);
445+
}
446+
441447
ret = Func(fcinfo);
448+
if constexpr (TArgsPolicy::VarArgs) {
449+
if (needToFree) {
450+
FreeVariadicArray(*fcinfo, batch.values.size());
451+
}
452+
}
453+
442454
if constexpr (IsFixedResult) {
443455
fixedResultData[i] = ui64(ret);
444456
fixedResultValidMask[i] = !fcinfo->isnull;

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1117,14 +1117,20 @@ class TPgResolvedCall : public TPgResolvedCallBase<TPgResolvedCall<UseContext>>
11171117
callInfo.args[i] = argDatum;
11181118
}
11191119

1120+
const bool needToFree = PrepareVariadicArray(callInfo, this->ProcDesc);
1121+
NUdf::TUnboxedValuePod res;
11201122
if constexpr (!UseContext) {
11211123
TPAllocScope call;
1122-
return this->DoCall(callInfo);
1124+
res = this->DoCall(callInfo);
1125+
} else {
1126+
res = this->DoCall(callInfo);
11231127
}
11241128

1125-
if constexpr (UseContext) {
1126-
return this->DoCall(callInfo);
1129+
if (needToFree) {
1130+
FreeVariadicArray(callInfo, this->ArgNodes.size());
11271131
}
1132+
1133+
return res;
11281134
}
11291135

11301136
private:
@@ -2618,6 +2624,7 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakePgKernel(TVector<TType*> argTy
26182624

26192625
state->flinfo.fn_expr = fmgrDataHolder->Build(procDesc);
26202626
state->FmgrDataHolder = fmgrDataHolder;
2627+
state->ProcDesc = &procDesc;
26212628

26222629
return arrow::Result(std::move(state));
26232630
};

‎ydb/library/yql/parser/pg_wrapper/ut/codegen_ut.cpp‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -158,12 +158,8 @@ Y_UNIT_TEST_SUITE(TPgCodegen) {
158158
kernelCtx.SetState(&state);
159159
FmgrInfo finfo;
160160
Zero(state.flinfo);
161-
if (fixed) {
162-
fmgr_info(NPg::LookupProc("date_eq", { 0, 0}).ProcId, &state.flinfo);
163-
} else {
164-
fmgr_info(NPg::LookupProc("textout", { 0} ).ProcId, &state.flinfo);
165-
}
166-
161+
state.ProcDesc = fixed ? &NPg::LookupProc("date_eq", { 0, 0 }) : &NPg::LookupProc("textout", { 0 });
162+
fmgr_info(state.ProcDesc->ProcId, &state.flinfo);
167163
state.context = nullptr;
168164
state.resultinfo = nullptr;
169165
state.fncollation = DEFAULT_COLLATION_OID;

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
#include <util/system/compiler.h>
2+
#include <ydb/library/yql/parser/pg_catalog/catalog.h>
3+
4+
extern "C" {
5+
#include "postgres.h"
6+
#include "fmgr.h"
7+
#include "utils/array.h"
8+
}
9+
10+
#undef Max
11+
12+
#include "utils.h"
13+
14+
namespace NYql {
15+
16+
void PrepareVariadicArraySlow(FunctionCallInfoBaseData& callInfo, const NPg::TProcDesc& desc) {
17+
const auto& elemDesc = NPg::LookupType(desc.VariadicType);
18+
19+
Datum varArgs[FUNC_MAX_ARGS];
20+
bool varArgsNulls[FUNC_MAX_ARGS];
21+
const ui32 nelems = callInfo.nargs - desc.ArgTypes.size();
22+
Y_ENSURE(nelems >= 1);
23+
for (ui32 i = desc.ArgTypes.size(); i < callInfo.nargs; ++i) {
24+
varArgs[i - desc.ArgTypes.size()] = callInfo.args[i].value;
25+
varArgsNulls[i - desc.ArgTypes.size()] = callInfo.args[i].isnull;
26+
}
27+
28+
callInfo.nargs = desc.ArgTypes.size() + 1;
29+
int dims[MAXDIM];
30+
int lbs[MAXDIM];
31+
dims[0] = nelems;
32+
lbs[0] = 1;
33+
34+
auto array = construct_md_array(varArgs, varArgsNulls, 1, dims, lbs,
35+
desc.VariadicType, elemDesc.TypeLen, elemDesc.PassByValue, elemDesc.TypeAlign);
36+
37+
auto& argDatum = callInfo.args[callInfo.nargs - 1];
38+
argDatum.value = PointerGetDatum(array);
39+
argDatum.isnull = false;
40+
}
41+
42+
void FreeVariadicArray(FunctionCallInfoBaseData& callInfo, ui32 originalArgs) {
43+
pfree(DatumGetPointer(callInfo.args[callInfo.nargs - 1].value));
44+
callInfo.nargs = originalArgs;
45+
}
46+
47+
}

‎ydb/library/yql/parser/pg_wrapper/utils.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,4 +137,16 @@ inline ui32 MakeTypeIOParam(const NPg::TTypeDesc& desc) {
137137
return desc.ElementTypeId ? desc.ElementTypeId : desc.TypeId;
138138
}
139139

140+
void PrepareVariadicArraySlow(FunctionCallInfoBaseData& callInfo, const NPg::TProcDesc& desc);
141+
void FreeVariadicArray(FunctionCallInfoBaseData& callInfo, ui32 originalArgs);
142+
143+
inline bool PrepareVariadicArray(FunctionCallInfoBaseData& callInfo, const NPg::TProcDesc& desc) {
144+
if (!desc.VariadicArgType || desc.VariadicArgType == desc.VariadicType) {
145+
return false;
146+
}
147+
148+
PrepareVariadicArraySlow(callInfo, desc);
149+
return true;
150+
}
151+
140152
}

‎ydb/library/yql/parser/pg_wrapper/ya.make‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ SRCS(
3838
cost_mocks.cpp
3939
syscache.cpp
4040
pg_utils_wrappers.cpp
41+
utils.cpp
4142
)
4243

4344
IF (ARCH_X86_64)

‎ydb/library/yql/tests/sql/dq_file/part12/canondata/result.json‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2572,6 +2572,28 @@
25722572
}
25732573
],
25742574
"test.test[pg-tpch-q16-default.txt-Results]": [],
2575+
"test.test[pg-variadic_array_arg-default.txt-Analyze]": [
2576+
{
2577+
"checksum": "b4dd508a329723c74293d80f0278c705",
2578+
"size": 505,
2579+
"uri": "https://{canondata_backend}/1942671/2b244e8340f9afab99b1136c03a3466d54265dea/resource.tar.gz#test.test_pg-variadic_array_arg-default.txt-Analyze_/plan.txt"
2580+
}
2581+
],
2582+
"test.test[pg-variadic_array_arg-default.txt-Debug]": [
2583+
{
2584+
"checksum": "22b71fe5b1ac2e9a705d2153308b4a05",
2585+
"size": 567,
2586+
"uri": "https://{canondata_backend}/1942671/2b244e8340f9afab99b1136c03a3466d54265dea/resource.tar.gz#test.test_pg-variadic_array_arg-default.txt-Debug_/opt.yql_patched"
2587+
}
2588+
],
2589+
"test.test[pg-variadic_array_arg-default.txt-Plan]": [
2590+
{
2591+
"checksum": "b4dd508a329723c74293d80f0278c705",
2592+
"size": 505,
2593+
"uri": "https://{canondata_backend}/1942671/2b244e8340f9afab99b1136c03a3466d54265dea/resource.tar.gz#test.test_pg-variadic_array_arg-default.txt-Plan_/plan.txt"
2594+
}
2595+
],
2596+
"test.test[pg-variadic_array_arg-default.txt-Results]": [],
25752597
"test.test[pg_catalog-pg_class-default.txt-Analyze]": [
25762598
{
25772599
"checksum": "c1f2d837c3623c81dd596a9877913fb8",

‎ydb/library/yql/tests/sql/hybrid_file/part2/canondata/result.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2071,6 +2071,20 @@
20712071
"uri": "https://{canondata_backend}/1936842/5461a7f0f4d722c81cba2eff5dd1d41bf3a77f80/resource.tar.gz#test.test_pg-tpch-q18-default.txt-Plan_/plan.txt"
20722072
}
20732073
],
2074+
"test.test[pg-variadic_array_arg-default.txt-Debug]": [
2075+
{
2076+
"checksum": "ffb1482578a6230fb1086e0c71c19735",
2077+
"size": 566,
2078+
"uri": "https://{canondata_backend}/1923547/be7e687ea36299e4a042c1495c58b793c69141af/resource.tar.gz#test.test_pg-variadic_array_arg-default.txt-Debug_/opt.yql_patched"
2079+
}
2080+
],
2081+
"test.test[pg-variadic_array_arg-default.txt-Plan]": [
2082+
{
2083+
"checksum": "b4dd508a329723c74293d80f0278c705",
2084+
"size": 505,
2085+
"uri": "https://{canondata_backend}/1923547/be7e687ea36299e4a042c1495c58b793c69141af/resource.tar.gz#test.test_pg-variadic_array_arg-default.txt-Plan_/plan.txt"
2086+
}
2087+
],
20742088
"test.test[pg_catalog-pg_description_pg_syntax-default.txt-Debug]": [
20752089
{
20762090
"checksum": "b7b94c65971f853c86abf32804ff3370",

‎ydb/library/yql/tests/sql/sql2yql/canondata/result.json‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13474,6 +13474,13 @@
1347413474
"uri": "https://{canondata_backend}/937458/e965a077f6aeb511184e419b0de866e505778b68/resource.tar.gz#test_sql2yql.test_pg-variadic_/sql.yql"
1347513475
}
1347613476
],
13477+
"test_sql2yql.test[pg-variadic_array_arg]": [
13478+
{
13479+
"checksum": "1815b2cb1ec493aefbcf189ebd533e2a",
13480+
"size": 608,
13481+
"uri": "https://{canondata_backend}/1936273/81aad5938509f97b7812a831a7a4698b6f82ebcc/resource.tar.gz#test_sql2yql.test_pg-variadic_array_arg_/sql.yql"
13482+
}
13483+
],
1347713484
"test_sql2yql.test[pg-wide_sort]": [
1347813485
{
1347913486
"checksum": "0a312c05be5c91545de12d451a697d00",

‎ydb/library/yql/tests/sql/suites/pg/variadic_array_arg.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
--!syntax_pg
2+
select json_extract_path('{"f2":{"f3":1},"f4":{"f5":99,"f6":"foo"}}', 'f4', 'f6')
3+

0 commit comments

Comments
 (0)