summaryrefslogtreecommitdiffstats
path: root/yql/essentials/parser/pg_wrapper/comp_factory.cpp
diff options
context:
space:
mode:
Diffstat (limited to 'yql/essentials/parser/pg_wrapper/comp_factory.cpp')
-rw-r--r--yql/essentials/parser/pg_wrapper/comp_factory.cpp99
1 files changed, 78 insertions, 21 deletions
diff --git a/yql/essentials/parser/pg_wrapper/comp_factory.cpp b/yql/essentials/parser/pg_wrapper/comp_factory.cpp
index 78b7e4c59fa..27601a078b3 100644
--- a/yql/essentials/parser/pg_wrapper/comp_factory.cpp
+++ b/yql/essentials/parser/pg_wrapper/comp_factory.cpp
@@ -463,6 +463,21 @@ public:
};
ApplyFillers(AllPgExtensionFillers, Y_ARRAY_SIZE(AllPgExtensionFillers), PgExtensionFillers_);
+ } else if (Table_ == "pg_collation") {
+ static const std::pair<const char*, TPgCollationFiller> AllPgCollationFillers[] = {
+ {"oid", [](const NPg::TCollationDesc& desc) { return ScalarDatumToPod(ObjectIdGetDatum(desc.Oid)); }},
+ {"collname", [](const NPg::TCollationDesc& desc) { return PointerDatumToPod((Datum)MakeFixedString(desc.Name, NAMEDATALEN)); }},
+ {"collnamespace", [](const NPg::TCollationDesc&) { return ScalarDatumToPod(ObjectIdGetDatum(PG_CATALOG_NAMESPACE)); }},
+ {"collowner", [](const NPg::TCollationDesc&) { return ScalarDatumToPod(ObjectIdGetDatum(1)); }},
+ {"collprovider", [](const NPg::TCollationDesc& desc) { return ScalarDatumToPod(CharGetDatum((char)desc.Provider)); }},
+ {"collisdeterministic", [](const NPg::TCollationDesc&) { return ScalarDatumToPod(BoolGetDatum(true)); }},
+ {"collencoding", [](const NPg::TCollationDesc& desc) { return ScalarDatumToPod(Int32GetDatum(desc.Encoding)); }},
+ {"collcollate", [](const NPg::TCollationDesc& desc) { return desc.Collate.empty() ? NUdf::TUnboxedValuePod() : PointerDatumToPod((Datum)MakeVar(desc.Collate)); }},
+ {"collctype", [](const NPg::TCollationDesc& desc) { return desc.Ctype.empty() ? NUdf::TUnboxedValuePod() : PointerDatumToPod((Datum)MakeVar(desc.Ctype)); }},
+ {"colliculocale", [](const NPg::TCollationDesc& desc) { return desc.IcuLocale.empty() ? NUdf::TUnboxedValuePod() : PointerDatumToPod((Datum)MakeVar(desc.IcuLocale)); }},
+ };
+
+ ApplyFillers(AllPgCollationFillers, Y_ARRAY_SIZE(AllPgCollationFillers), PgCollationFillers_);
}
} else {
if (Table_ == "tables") {
@@ -864,6 +879,19 @@ public:
sysFiller.Fill(items);
rows.emplace_back(row);
});
+ } else if (Table_ == "pg_collation") {
+ NPg::EnumCollation([&](ui32, const NPg::TCollationDesc& desc) {
+ NUdf::TUnboxedValue* items;
+ auto row = compCtx.HolderFactory.CreateDirectArrayHolder(PgCollationFillers_.size(), items);
+ for (ui32 i = 0; i < PgCollationFillers_.size(); ++i) {
+ if (PgCollationFillers_[i]) {
+ items[i] = PgCollationFillers_[i](desc);
+ }
+ }
+
+ sysFiller.Fill(items);
+ rows.emplace_back(row);
+ });
}
} else {
if (Table_ == "tables") {
@@ -952,6 +980,9 @@ private:
using TPgProcFiller = NUdf::TUnboxedValuePod(*)(const NPg::TProcDesc&);
TVector<TPgProcFiller> PgProcFillers_;
+ using TPgCollationFiller = NUdf::TUnboxedValuePod(*)(const NPg::TCollationDesc&);
+ TVector<TPgCollationFiller> PgCollationFillers_;
+
using TPgAggregateFiller = NUdf::TUnboxedValuePod(*)(const NPg::TAggregateDesc&);
TVector<TPgAggregateFiller> PgAggregateFillers_;
@@ -1033,8 +1064,8 @@ private:
};
struct TPgResolvedCallNodeState : public TPgResolvedCallState, public TComputationValue<TPgResolvedCallNodeState> {
- TPgResolvedCallNodeState(TMemoryUsageInfo* memInfo, ui32 numArgs, const FmgrInfo* finfo)
- : TPgResolvedCallState(numArgs, finfo)
+ TPgResolvedCallNodeState(TMemoryUsageInfo* memInfo, ui32 numArgs, const FmgrInfo* finfo, ui32 collationOid)
+ : TPgResolvedCallState(numArgs, finfo, collationOid)
, TComputationValue(memInfo) {}
};
@@ -1043,10 +1074,11 @@ class TPgResolvedCallNode : public TPgResolvedCall<UseContext>, public TMutableC
typedef TMutableComputationNode<TPgResolvedCallNode<UseContext>> TBaseComputation;
public:
TPgResolvedCallNode(TComputationMutables& mutables, const std::string_view& name, ui32 id,
- TComputationNodePtrVector&& argNodes, TVector<TType*>&& argTypes, TType* returnType)
+ TComputationNodePtrVector&& argNodes, TVector<TType*>&& argTypes, TType* returnType, ui32 collationOid)
: TPgResolvedCall<UseContext>(name, id, std::move(argTypes), returnType)
, TBaseComputation(mutables)
, ArgNodes(std::move(argNodes))
+ , CollationOid(collationOid)
, StateIndex(mutables.CurValueIndex++)
{
Y_ENSURE(ArgNodes.size() == this->ArgDesc.size());
@@ -1068,13 +1100,14 @@ private:
TPgResolvedCallNodeState& GetState(TComputationContext& compCtx) const {
auto& result = compCtx.MutableValues[this->StateIndex];
if (!result.HasValue()) {
- result = compCtx.HolderFactory.Create<TPgResolvedCallNodeState>(this->ArgNodes.size(), &this->FInfo);
+ result = compCtx.HolderFactory.Create<TPgResolvedCallNodeState>(this->ArgNodes.size(), &this->FInfo, CollationOid);
}
return *static_cast<TPgResolvedCallNodeState*>(result.AsBoxed().Get());
}
const TComputationNodePtrVector ArgNodes;
+ const ui32 CollationOid;
const ui32 StateIndex;
};
@@ -1087,14 +1120,14 @@ private:
public:
TIterator(TMemoryUsageInfo* memInfo, const std::string_view& name, const TUnboxedValueVector& args,
const TVector<NPg::TTypeDesc>& argDesc, const NPg::TTypeDesc& retTypeDesc, const NPg::TProcDesc& procDesc,
- const FmgrInfo* fInfo, const TStructType* structType, const TVector<NPg::TTypeDesc>& structTypeDesc, const THolderFactory& holderFactory)
+ const FmgrInfo* fInfo, ui32 collationOid, const TStructType* structType, const TVector<NPg::TTypeDesc>& structTypeDesc, const THolderFactory& holderFactory)
: TComputationValue<TIterator>(memInfo)
, Name(name)
, Args(args)
, ArgDesc(argDesc)
, RetTypeDesc(retTypeDesc)
, ProcDesc(procDesc)
- , CallInfo(argDesc.size(), fInfo)
+ , CallInfo(argDesc.size(), fInfo, collationOid)
, StructType(structType)
, StructTypeDesc(structTypeDesc)
, HolderFactory(holderFactory)
@@ -1302,7 +1335,7 @@ private:
TListValue(TMemoryUsageInfo* memInfo, TComputationContext& compCtx,
const std::string_view& name, TUnboxedValueVector&& args, const TVector<NPg::TTypeDesc>& argDesc,
- const NPg::TTypeDesc& retTypeDesc, const NPg::TProcDesc& procDesc, const FmgrInfo* fInfo,
+ const NPg::TTypeDesc& retTypeDesc, const NPg::TProcDesc& procDesc, const FmgrInfo* fInfo, ui32 collationOid,
const TStructType* structType, const TVector<NPg::TTypeDesc>& structTypeDesc, const THolderFactory& /*holderFactory*/)
: TCustomListValue(memInfo)
, CompCtx(compCtx)
@@ -1312,6 +1345,7 @@ private:
, RetTypeDesc(retTypeDesc)
, ProcDesc(procDesc)
, FInfo(fInfo)
+ , CollationOid(collationOid)
, StructType(structType)
, StructTypeDesc(structTypeDesc)
{
@@ -1319,7 +1353,7 @@ private:
private:
NUdf::TUnboxedValue GetListIterator() const final {
- return CompCtx.HolderFactory.Create<TIterator>(Name, Args, ArgDesc, RetTypeDesc, ProcDesc, FInfo, StructType, StructTypeDesc, CompCtx.HolderFactory);
+ return CompCtx.HolderFactory.Create<TIterator>(Name, Args, ArgDesc, RetTypeDesc, ProcDesc, FInfo, CollationOid, StructType, StructTypeDesc, CompCtx.HolderFactory);
}
TComputationContext& CompCtx;
@@ -1329,16 +1363,18 @@ private:
const NPg::TTypeDesc& RetTypeDesc;
const NPg::TProcDesc& ProcDesc;
const FmgrInfo* FInfo;
+ const ui32 CollationOid;
const TStructType* StructType;
const TVector<NPg::TTypeDesc>& StructTypeDesc;
};
public:
TPgResolvedMultiCall(TComputationMutables& mutables, const std::string_view& name, ui32 id,
- TComputationNodePtrVector&& argNodes, TVector<TType*>&& argTypes, TType* returnType, const TStructType* structType)
+ TComputationNodePtrVector&& argNodes, TVector<TType*>&& argTypes, TType* returnType, const TStructType* structType, ui32 collationOid)
: TPgResolvedCallBase(name, id, std::move(argTypes), returnType, true, structType)
, TBaseComputation(mutables)
, ArgNodes(std::move(argNodes))
+ , CollationOid(collationOid)
, StructType(structType)
{
Y_ENSURE(ArgNodes.size() == ArgDesc.size());
@@ -1352,7 +1388,7 @@ public:
args.push_back(value);
}
- return compCtx.HolderFactory.Create<TListValue>(compCtx, Name, std::move(args), ArgDesc, RetTypeDesc, ProcDesc, &FInfo, StructType, StructTypeDesc, compCtx.HolderFactory);
+ return compCtx.HolderFactory.Create<TListValue>(compCtx, Name, std::move(args), ArgDesc, RetTypeDesc, ProcDesc, &FInfo, CollationOid, StructType, StructTypeDesc, compCtx.HolderFactory);
}
private:
@@ -1362,6 +1398,7 @@ private:
}
}
const TComputationNodePtrVector ArgNodes;
+ const ui32 CollationOid;
const TStructType* StructType;
};
@@ -2757,7 +2794,7 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakeToPgKernel(TType* inputType, T
return kernel;
}
-std::shared_ptr<arrow::compute::ScalarKernel> MakePgKernel(TVector<TType*> argTypes, TType* resultType, TExecFunc execFunc, ui32 procId) {
+std::shared_ptr<arrow::compute::ScalarKernel> MakePgKernel(TVector<TType*> argTypes, TType* resultType, TExecFunc execFunc, ui32 procId, ui32 collationOid) {
std::shared_ptr<arrow::DataType> returnArrowType;
MKQL_ENSURE(ConvertArrowType(AS_TYPE(TBlockType, resultType)->GetItemType(), returnArrowType), "Unsupported arrow type");
auto kernel = std::make_shared<arrow::compute::ScalarKernel>(ConvertToInputTypes(argTypes), ConvertToOutputType(resultType),
@@ -2778,14 +2815,14 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakePgKernel(TVector<TType*> argTy
}
kernel->null_handling = arrow::compute::NullHandling::COMPUTED_NO_PREALLOCATE;
- kernel->init = [procId, pgArgTypes](arrow::compute::KernelContext*, const arrow::compute::KernelInitArgs&) {
+ kernel->init = [procId, pgArgTypes, collationOid](arrow::compute::KernelContext*, const arrow::compute::KernelInitArgs&) {
auto state = std::make_unique<TPgKernelState>();
Zero(state->flinfo);
GetPgFuncAddr(procId, state->flinfo);
YQL_ENSURE(state->flinfo.fn_addr);
state->resultinfo = nullptr;
state->context = nullptr;
- state->fncollation = DEFAULT_COLLATION_OID;
+ state->fncollation = collationOid ? collationOid : DEFAULT_COLLATION_OID;
const auto& procDesc = NPg::LookupProc(procId);
const auto& retTypeDesc = NPg::LookupType(procDesc.ResultType);
state->Name = procDesc.Name;
@@ -2855,7 +2892,8 @@ TComputationNodeFactory GetPgFactory() {
return new TPgToRecord(ctx.Mutables, input, structType, std::move(members));
}
- if (name == "PgResolvedCall") {
+ if (name == "PgResolvedCall" || name == "PgResolvedCall2") {
+ const bool hasCollation = (name == "PgResolvedCall2");
const auto useContextData = AS_VALUE(TDataLiteral, callable.GetInput(0));
const auto rangeFunctionData = AS_VALUE(TDataLiteral, callable.GetInput(1));
const auto nameData = AS_VALUE(TDataLiteral, callable.GetInput(2));
@@ -2864,9 +2902,18 @@ TComputationNodeFactory GetPgFactory() {
auto rangeFunction = rangeFunctionData->AsValue().Get<bool>();
auto name = nameData->AsValue().AsStringRef();
auto id = idData->AsValue().Get<ui32>();
+
+ ui32 collationOid = 0;
+ ui32 argsStart = 4;
+ if (hasCollation) {
+ const auto collationOidData = AS_VALUE(TDataLiteral, callable.GetInput(4));
+ collationOid = collationOidData->AsValue().Get<ui32>();
+ argsStart = 5;
+ }
+
TComputationNodePtrVector argNodes;
TVector<TType*> argTypes;
- for (ui32 i = 4; i < callable.GetInputsCount(); ++i) {
+ for (ui32 i = argsStart; i < callable.GetInputsCount(); ++i) {
argNodes.emplace_back(LocateNode(ctx.NodeLocator, callable, i));
argTypes.emplace_back(callable.GetInput(i).GetStaticType());
}
@@ -2883,25 +2930,35 @@ TComputationNodeFactory GetPgFactory() {
if (isList) {
YQL_ENSURE(!useContext);
- return new TPgResolvedMultiCall(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), itemType, structType);
+ return new TPgResolvedMultiCall(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), itemType, structType, collationOid);
} else {
YQL_ENSURE(!structType);
if (useContext) {
- return new TPgResolvedCallNode<true>(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), returnType);
+ return new TPgResolvedCallNode<true>(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), returnType, collationOid);
} else {
- return new TPgResolvedCallNode<false>(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), returnType);
+ return new TPgResolvedCallNode<false>(ctx.Mutables, name, id, std::move(argNodes), std::move(argTypes), returnType, collationOid);
}
}
}
- if (name == "BlockPgResolvedCall") {
+ if (name == "BlockPgResolvedCall" || name == "BlockPgResolvedCall2") {
+ const bool hasCollation = (name == "BlockPgResolvedCall2");
const auto nameData = AS_VALUE(TDataLiteral, callable.GetInput(0));
const auto idData = AS_VALUE(TDataLiteral, callable.GetInput(1));
auto name = nameData->AsValue().AsStringRef();
auto id = idData->AsValue().Get<ui32>();
+
+ ui32 collationOid = 0;
+ ui32 argsStart = 2;
+ if (hasCollation) {
+ const auto collationOidData = AS_VALUE(TDataLiteral, callable.GetInput(2));
+ collationOid = collationOidData->AsValue().Get<ui32>();
+ argsStart = 3;
+ }
+
TComputationNodePtrVector argNodes;
TVector<TType*> argTypes;
- for (ui32 i = 2; i < callable.GetInputsCount(); ++i) {
+ for (ui32 i = argsStart; i < callable.GetInputsCount(); ++i) {
argNodes.emplace_back(LocateNode(ctx.NodeLocator, callable, i));
argTypes.emplace_back(callable.GetInput(i).GetStaticType());
}
@@ -2909,7 +2966,7 @@ TComputationNodeFactory GetPgFactory() {
auto returnType = callable.GetType()->GetReturnType();
auto execFunc = FindExec(id);
YQL_ENSURE(execFunc);
- auto kernel = MakePgKernel(argTypes, returnType, execFunc, id);
+ auto kernel = MakePgKernel(argTypes, returnType, execFunc, id, collationOid);
return new TBlockFuncNode(ctx.Mutables, ctx.RuntimeSettings->DatumValidation.Get(), callable.GetType()->GetName(), std::move(argNodes), argTypes, returnType, *kernel, kernel);
}
generated by cgit v1.3 (git 2.53.0) at 2026-07-19 00:35:45 +0000