From 69259d92949b108e289d527f8692af80771abd84 Mon Sep 17 00:00:00 2001
From: "zhenshan.cao" <zhenshan.cao@zilliz.com>
Date: Wed, 8 Dec 2021 15:53:05 +0800
Subject: [PATCH] Call ReleaseUnused of memorypool in arrow (#12893)

Signed-off-by: zhenshan.cao <zhenshan.cao@zilliz.com>
---
 internal/storage/cwrapper/ParquetWrapper.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/internal/storage/cwrapper/ParquetWrapper.cpp b/internal/storage/cwrapper/ParquetWrapper.cpp
index 2d30ec8e0..9b32e513b 100644
--- a/internal/storage/cwrapper/ParquetWrapper.cpp
+++ b/internal/storage/cwrapper/ParquetWrapper.cpp
@@ -307,7 +307,8 @@ CStatus FinishPayloadWriter(CPayloadWriter payloadWriter) {
     }
     auto table = arrow::Table::Make(p->schema, {array});
     p->output = std::make_shared<wrapper::PayloadOutputStream>();
-    ast = parquet::arrow::WriteTable(*table, arrow::default_memory_pool(), p->output, 1024 * 1024 * 1024);
+    auto mem_pool = arrow::default_memory_pool();
+    ast = parquet::arrow::WriteTable(*table, mem_pool, p->output, 1024 * 1024 * 1024);
     if (!ast.ok()) {
       st.error_code = static_cast<int>(ErrorCode::UNEXPECTED_ERROR);
       st.error_msg = ErrorMsg(ast.message());
@@ -344,6 +345,8 @@ CStatus ReleasePayloadWriter(CPayloadWriter handler) {
   st.error_msg = nullptr;
   auto p = reinterpret_cast<wrapper::PayloadWriter *>(handler);
   if (p != nullptr) delete p;
+  auto mem_pool = arrow::default_memory_pool();
+  mem_pool->ReleaseUnused();
   return st;
 }
 
@@ -352,7 +355,8 @@ CPayloadReader NewPayloadReader(int columnType, uint8_t *buffer, int64_t buf_siz
   auto p = new wrapper::PayloadReader;
   p->bValues = nullptr;
   p->input = std::make_shared<wrapper::PayloadInputStream>(buffer, buf_size);
-  auto st = parquet::arrow::OpenFile(p->input, arrow::default_memory_pool(), &p->reader);
+  auto mem_pool = arrow::default_memory_pool();
+  auto st = parquet::arrow::OpenFile(p->input, mem_pool, &p->reader);
   if (!st.ok()) {
     delete p;
     return nullptr;
@@ -366,7 +370,6 @@ CPayloadReader NewPayloadReader(int columnType, uint8_t *buffer, int64_t buf_siz
   assert(p->column != nullptr);
   assert(p->column->chunks().size() == 1);
   p->array = p->column->chunk(0);
-
   switch (columnType) {
     case ColumnType::BOOL :
     case ColumnType::INT8 :
@@ -533,5 +536,7 @@ CStatus ReleasePayloadReader(CPayloadReader payloadReader) {
   auto p = reinterpret_cast<wrapper::PayloadReader *>(payloadReader);
   delete[] p->bValues;
   delete p;
+  auto mem_pool = arrow::default_memory_pool();
+  mem_pool->ReleaseUnused();
   return st;
 }
-- 
GitLab