fix

Author: nikvas0

ydb/core/kqp/runtime/kqp_write_actor.cpp

@@ -137,6 +137,7 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
         , Callbacks(args.Callback)
         , Counters(counters)
         , TypeEnv(args.TypeEnv)
+        , Alloc(args.Alloc)
         , TxId(args.TxId)
         , TableId(
             Settings.GetTable().GetOwnerId(),
@@ -693,7 +694,8 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
                             : kMaxBatchesPerMessage),
                     },
                     std::move(columnsMetadata),
-                    TypeEnv);
+                    TypeEnv,
+                    Alloc);
             } catch (...) {
                 RuntimeError(
                     CurrentExceptionMessage(),
@@ -730,6 +732,7 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
     NYql::NDq::IDqComputeActorAsyncOutput::ICallbacks * Callbacks = nullptr;
     TIntrusivePtr<TKqpCounters> Counters;
     const NMiniKQL::TTypeEnvironment& TypeEnv;
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
 
     const NYql::NDq::TTxId TxId;
     const TTableId TableId;

ydb/core/kqp/runtime/kqp_write_table.cpp

@@ -1184,14 +1184,17 @@ class TShardedWriteController : public IShardedWriteController {
     TShardedWriteController(
         const TShardedWriteControllerSettings settings,
         TVector<NKikimrKqp::TKqpColumnMetadataProto>&& inputColumnsMetadata,
-        const NMiniKQL::TTypeEnvironment& typeEnv)
+        const NMiniKQL::TTypeEnvironment& typeEnv,
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc)
         : Settings(settings)
         , InputColumnsMetadata(std::move(inputColumnsMetadata))
-        , TypeEnv(typeEnv) {
+        , TypeEnv(typeEnv)
+        , Alloc(alloc) {
     }
 
     ~TShardedWriteController() {
-        auto allocGuard = TypeEnv.BindAllocator();
+        Y_ABORT_UNLESS(Alloc);
+        TGuard<NMiniKQL::TScopedAlloc> allocGuard(*Alloc);
         ShardsInfo.Clear();
         Serializer = nullptr;
     }
@@ -1237,6 +1240,7 @@ class TShardedWriteController : public IShardedWriteController {
     TShardedWriteControllerSettings Settings;
     TVector<NKikimrKqp::TKqpColumnMetadataProto> InputColumnsMetadata;
     const NMiniKQL::TTypeEnvironment& TypeEnv;
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
 
     TShardsInfo ShardsInfo;
     bool Closed = false;
@@ -1250,8 +1254,10 @@ class TShardedWriteController : public IShardedWriteController {
 IShardedWriteControllerPtr CreateShardedWriteController(
         const TShardedWriteControllerSettings& settings,
         TVector<NKikimrKqp::TKqpColumnMetadataProto>&& inputColumns,
-        const NMiniKQL::TTypeEnvironment& typeEnv) {
-    return MakeIntrusive<TShardedWriteController>(settings, std::move(inputColumns), typeEnv);
+        const NMiniKQL::TTypeEnvironment& typeEnv,
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc) {
+    return MakeIntrusive<TShardedWriteController>(
+        settings, std::move(inputColumns), typeEnv, alloc);
 }
 
 }

ydb/core/kqp/runtime/kqp_write_table.h

@@ -64,7 +64,8 @@ struct TShardedWriteControllerSettings {
 IShardedWriteControllerPtr CreateShardedWriteController(
     const TShardedWriteControllerSettings& settings,
     TVector<NKikimrKqp::TKqpColumnMetadataProto>&& inputColumns,
-    const NMiniKQL::TTypeEnvironment& typeEnv);
+    const NMiniKQL::TTypeEnvironment& typeEnv,
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc);
 
 }
 }

ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io.h

@@ -299,6 +299,7 @@ struct IDqAsyncIoFactory : public TThrRefBase {
         const THashMap<TString, TString>& TaskParams;
         const NKikimr::NMiniKQL::TTypeEnvironment& TypeEnv;
         const NKikimr::NMiniKQL::THolderFactory& HolderFactory;
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
         IRandomProvider *const RandomProvider;
     };
 

ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h

@@ -1359,6 +1359,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
                         .TaskParams = taskParams,
                         .TypeEnv = typeEnv,
                         .HolderFactory = holderFactory,
+                        .Alloc = Alloc,
                         .RandomProvider = randomProvider
                     });
             } catch (const std::exception& ex) {