Memory tracking for sinks (#16102)

Author: nikvas0

.github/config/muted_ya.txt

@@ -25,8 +25,6 @@ ydb/core/kqp/ut/olap KqpOlapWrite.TierDraftsGCWithRestart
 ydb/core/kqp/ut/olap [*/*] chunk chunk
 ydb/core/kqp/ut/query KqpAnalyze.AnalyzeTable+ColumnStore
 ydb/core/kqp/ut/query KqpAnalyze.AnalyzeTable-ColumnStore
-ydb/core/kqp/ut/query KqpLimits.OutOfSpaceYQLUpsertFail+useSink
-ydb/core/kqp/ut/query KqpLimits.QSReplySizeEnsureMemoryLimits+useSink
 ydb/core/kqp/ut/query KqpStats.SysViewClientLost
 ydb/core/kqp/ut/scheme KqpOlapScheme.TenThousandColumns
 ydb/core/kqp/ut/scheme KqpScheme.AlterAsyncReplication

ydb/core/formats/arrow/arrow_batch_builder.cpp

@@ -183,9 +183,13 @@ bool TRecordBatchReader::DeserializeFromStrings(const TString& schemaString, con
     return true;
 }
 
-TArrowBatchBuilder::TArrowBatchBuilder(arrow::Compression::type codec, const std::set<std::string>& notNullColumns)
+TArrowBatchBuilder::TArrowBatchBuilder(
+        arrow::Compression::type codec,
+        const std::set<std::string>& notNullColumns,
+        arrow::MemoryPool* memoryPool)
     : WriteOptions(arrow::ipc::IpcWriteOptions::Defaults())
     , NotNullColumns(notNullColumns)
+    , MemoryPool(memoryPool)
 {
     Y_ABORT_UNLESS(arrow::util::Codec::IsAvailable(codec));
     auto resCodec = arrow::util::Codec::Create(codec);
@@ -201,7 +205,7 @@ arrow::Status TArrowBatchBuilder::Start(const std::vector<std::pair<TString, NSc
     if (!schema.ok()) {
         return arrow::Status::FromArgs(schema.status().code(), "Cannot make arrow schema: ", schema.status().ToString());
     }
-    auto status = arrow::RecordBatchBuilder::Make(*schema, arrow::default_memory_pool(), RowsToReserve, &BatchBuilder);
+    auto status = arrow::RecordBatchBuilder::Make(*schema, MemoryPool, RowsToReserve, &BatchBuilder);
     NumRows = NumBytes = 0;
     if (!status.ok()) {
         return arrow::Status::FromArgs(schema.status().code(), "Cannot make arrow builder: ", status.ToString());

ydb/core/formats/arrow/arrow_batch_builder.h

@@ -149,7 +149,10 @@ class TArrowBatchBuilder : public NKikimr::IBlockBuilder {
 
     /// @note compression is disabled by default KIKIMR-11690
     // Allowed codecs: UNCOMPRESSED, LZ4_FRAME, ZSTD
-    TArrowBatchBuilder(arrow::Compression::type codec = arrow::Compression::UNCOMPRESSED, const std::set<std::string>& notNullColumns = {});
+    TArrowBatchBuilder(
+        arrow::Compression::type codec = arrow::Compression::UNCOMPRESSED,
+        const std::set<std::string>& notNullColumns = {},
+        arrow::MemoryPool* memoryPool = arrow::default_memory_pool());
     ~TArrowBatchBuilder() = default;
 
     bool Start(const std::vector<std::pair<TString, NScheme::TTypeInfo>>& columns,
@@ -197,6 +200,7 @@ class TArrowBatchBuilder : public NKikimr::IBlockBuilder {
     std::shared_ptr<arrow::RecordBatch> Batch;
     size_t RowsToReserve{DEFAULT_ROWS_TO_RESERVE};
     const std::set<std::string> NotNullColumns;
+    arrow::MemoryPool* MemoryPool;
 
 protected:
     size_t NumRows{0};

ydb/core/kqp/common/buffer/buffer.h

@@ -12,6 +12,7 @@ struct TKqpBufferWriterSettings {
     NWilson::TTraceId TraceId;
     TIntrusivePtr<TKqpCounters> Counters;
     TIntrusivePtr<NTxProxy::TTxProxyMon> TxProxyMon;
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
 };
 
 NActors::IActor* CreateKqpBufferWriterActor(TKqpBufferWriterSettings&& settings);