Fix usage single UDF object from many threads in DQ. (#7436)

Author: Tony-Romanov

ydb/core/tx/datashard/datashard_kqp_lookup_table.cpp

@@ -121,7 +121,7 @@ class TKqpLookupRowsWrapper : public TStatelessFlowComputationNode<TKqpLookupRow
             switch (keysValues.Fetch(key)) {
                 case NUdf::EFetchStatus::Ok: {
                     TVector<TCell> keyCells(ParseResult.KeyIndices.size());
-                    FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, keyCells, *ctx.TypeEnv);
+                    FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, keyCells, ctx.TypeEnv);
 
                     NUdf::TUnboxedValue result;
                     TKqpTableStats stats;
@@ -203,10 +203,10 @@ class TKqpLookupTableWrapper : public TStatelessFlowComputationNode<TKqpLookupTa
                         MKQL_ENSURE_S(tableInfo);
 
                         TVector<TCell> fromCells(tableInfo->KeyColumns.size());
-                        FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, fromCells, *ctx.TypeEnv);
+                        FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, fromCells, ctx.TypeEnv);
 
                         TVector<TCell> toCells(ParseResult.KeyIndices.size());
-                        FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, toCells, *ctx.TypeEnv);
+                        FillKeyTupleValue(key, ParseResult.KeyIndices, ParseResult.KeyTypes, toCells, ctx.TypeEnv);
 
                         auto range = TTableRange(fromCells, true, toCells, true);
 

ydb/core/tx/datashard/datashard_kqp_read_table.cpp

@@ -269,10 +269,10 @@ class TKqpWideReadTableWrapper : public TKqpWideReadTableWrapperBase<IsReverse>
     EFetchResult ReadValue(TComputationContext& ctx, NUdf::TUnboxedValue* const* output) const final {
         if (!this->Iterator) {
             TVector<TCell> fromCells;
-            BuildKeyTupleCells(ParseResult.FromTuple->GetType(), FromNode->GetValue(ctx), fromCells, *ctx.TypeEnv);
+            BuildKeyTupleCells(ParseResult.FromTuple->GetType(), FromNode->GetValue(ctx), fromCells, ctx.TypeEnv);
 
             TVector<TCell> toCells;
-            BuildKeyTupleCells(ParseResult.ToTuple->GetType(), ToNode->GetValue(ctx), toCells, *ctx.TypeEnv);
+            BuildKeyTupleCells(ParseResult.ToTuple->GetType(), ToNode->GetValue(ctx), toCells, ctx.TypeEnv);
 
             auto range = TTableRange(fromCells, ParseResult.FromInclusive, toCells, ParseResult.ToInclusive);
 
@@ -328,7 +328,7 @@ class TKqpWideReadTableRangesWrapper : public TKqpWideReadTableWrapperBase<IsRev
         if (!RangeId) {
             const auto localTid = this->ComputeCtx.GetLocalTableId(ParseResult.TableId);
             const auto* tableInfo = this->ComputeCtx.Database->GetScheme().GetTableInfo(localTid);
-            Ranges = CreateTableRanges<IsReverse>(ParseResult, RangesNode, *ctx.TypeEnv, ctx, tableInfo->KeyColumns.size());
+            Ranges = CreateTableRanges<IsReverse>(ParseResult, RangesNode, ctx.TypeEnv, ctx, tableInfo->KeyColumns.size());
             RangeId = 0;
 
             if (ItemsLimit) {

ydb/core/tx/program/registry.cpp

@@ -7,18 +7,18 @@
 
 namespace NKikimr::NOlap {
 
-::NTls::TValue<NMiniKQL::IBuiltinFunctionRegistry::TPtr> Registry;
+::NTls::TValue<TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry>> Registry;
 
 bool TKernelsRegistry::Parse(const TString& serialized) {
     Y_ABORT_UNLESS(!!serialized);
     if (!Registry.Get()) {
-        Registry = NMiniKQL::CreateBuiltinRegistry();
+        auto registry = NMiniKQL::CreateFunctionRegistry(NMiniKQL::CreateBuiltinRegistry())->Clone();
+        NMiniKQL::FillStaticModules(*registry.Get());
+        Registry = std::move(registry);
     }
-    auto copy = Registry.Get();
-    auto functionRegistry = NMiniKQL::CreateFunctionRegistry(std::move(copy))->Clone();
-    NMiniKQL::FillStaticModules(*functionRegistry);
+
     auto nodeFactory = NMiniKQL::GetBuiltinFactory();
-    auto kernels =  NYql::LoadKernels(serialized, *functionRegistry, nodeFactory);
+    auto kernels =  NYql::LoadKernels(serialized, *Registry.Get(), nodeFactory);
     Kernels.swap(kernels);
     for (const auto& kernel : Kernels) {
         arrow::compute::Arity arity(kernel->signature->in_types().size(), kernel->signature->is_varargs());

ydb/library/yql/dq/runtime/dq_tasks_runner.cpp

@@ -457,7 +457,7 @@ class TDqTaskRunner : public IDqTaskRunner {
         auto opts = CreatePatternOpts(task, Alloc(), TypeEnv());
 
         AllocatedHolder->ProgramParsed.CompGraph = AllocatedHolder->ProgramParsed.GetPattern()->Clone(
-            opts.ToComputationOptions(*Context.RandomProvider, *Context.TimeProvider, &TypeEnv()));
+            opts.ToComputationOptions(*Context.RandomProvider, *Context.TimeProvider));
 
         TBindTerminator term(AllocatedHolder->ProgramParsed.CompGraph->GetTerminator());
 

ydb/library/yql/minikql/comp_nodes/mkql_scalar_apply.cpp

@@ -44,15 +44,17 @@ class TScalarApplyWrapper : public TMutableComputationNode<TScalarApplyWrapper>
     struct TKernelState : public arrow::compute::KernelState {
         TKernelState(const TVector<TType*>& argsTypes, TType* returnType, const TComputationContext& originalContext)
             : Alloc(__LOCATION__)
+            , TypeEnv(Alloc)
             , MemInfo("ScalarApply")
-            , HolderFactory(Alloc.Ref(), MemInfo)
+            , FunctionRegistry(originalContext.HolderFactory.GetFunctionRegistry()->Clone())
+            , HolderFactory(Alloc.Ref(), MemInfo, FunctionRegistry.Get())
             , ValueBuilder(HolderFactory, NUdf::EValidatePolicy::Exception)
             , PgBuilder(NYql::CreatePgBuilder())
             , Accessors(argsTypes, returnType, *PgBuilder)
             , RandomProvider(CreateDefaultRandomProvider())
             , TimeProvider(CreateDefaultTimeProvider())
             , Ctx(HolderFactory, &ValueBuilder, TComputationOptsFull(
-                nullptr, Alloc.Ref(), *RandomProvider, *TimeProvider, NUdf::EValidatePolicy::Exception, nullptr),
+                nullptr, Alloc.Ref(), TypeEnv, *RandomProvider, *TimeProvider, NUdf::EValidatePolicy::Exception, originalContext.SecureParamsProvider, originalContext.CountersProvider),
                 originalContext.Mutables, *NYql::NUdf::GetYqlMemoryPool())
         {
             Alloc.Ref().EnableArrowTracking = false;
@@ -65,7 +67,9 @@ class TScalarApplyWrapper : public TMutableComputationNode<TScalarApplyWrapper>
         }
 
         TScopedAlloc Alloc;
+        TTypeEnvironment TypeEnv;
         TMemoryUsageInfo MemInfo;
+        const IFunctionRegistry::TPtr FunctionRegistry;
         THolderFactory HolderFactory;
         TDefaultValueBuilder ValueBuilder;
         std::unique_ptr<NUdf::IPgBuilder> PgBuilder;