Skip to content

Fix memory pool lifetime for Arrow stream buffers #464

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mrdrivingduck wants to merge 1 commit into
base: main
Choose a base branch
from
+136 −2
Open

Fix memory pool lifetime for Arrow stream buffers #464

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions src/paimon/common/utils/arrow/arrow_input_stream_adapter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,19 @@ arrow::Status ValidateArrowIoRange(int64_t value, const char* name) {
return arrow::Status::OK();
}

struct BufferWithMemoryPool {

@zjw1111 zjw1111 Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have a performance concern, as this would increment the shared_ptr reference count on every I/O operation. Have you evaluated the performance overhead of introducing shared_ptr here?

@mrdrivingduck mrdrivingduck Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran a pure in-memory microbenchmark. Results (before -> after, ns/op):

Read size Read ReadAt ReadAsync
4 KiB, 1 thread +10.4% +10.9% +7.9%
64 KiB, 1 thread +3.7% +8.4% +4.6%
1 MiB, 1 thread +0.0% -1.2% +0.1%

The overhead becomes smaller as the read size increases. The fix is still needed to keep the memory pool alive while returned buffers are referenced; otherwise, the ownership between the pool, reader, and buffers would need to change.

std::shared_ptr<arrow::MemoryPool> pool;
std::shared_ptr<arrow::Buffer> buffer;
};

std::shared_ptr<arrow::Buffer> KeepMemoryPoolAlive(std::shared_ptr<arrow::Buffer> buffer,
const std::shared_ptr<arrow::MemoryPool>& pool) {
auto holder =
std::make_shared<BufferWithMemoryPool>(BufferWithMemoryPool{pool, std::move(buffer)});
auto* buffer_ptr = holder->buffer.get();
return std::shared_ptr<arrow::Buffer>(std::move(holder), buffer_ptr);
}
Comment thread
mrdrivingduck marked this conversation as resolved.

} // namespace

ArrowInputStreamAdapter::ArrowInputStreamAdapter(
Expand Down Expand Up @@ -79,7 +92,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::Read(int6
if (read_bytes < nbytes) {
ARROW_RETURN_NOT_OK(buffer->Resize(read_bytes));
}
return std::shared_ptr<arrow::Buffer>(std::move(buffer));
return KeepMemoryPoolAlive(std::shared_ptr<arrow::Buffer>(std::move(buffer)), pool_);
}

arrow::Result<int64_t> ArrowInputStreamAdapter::ReadAt(int64_t position, int64_t nbytes,
Expand All @@ -104,7 +117,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::ReadAt(in
if (read_bytes < nbytes) {
ARROW_RETURN_NOT_OK(buffer->Resize(read_bytes));
}
return std::shared_ptr<arrow::Buffer>(std::move(buffer));
return KeepMemoryPoolAlive(std::shared_ptr<arrow::Buffer>(std::move(buffer)), pool_);
}

arrow::Future<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::ReadAsync(
Expand All @@ -128,6 +141,7 @@ arrow::Future<std::shared_ptr<arrow::Buffer>> ArrowInputStreamAdapter::ReadAsync
return fut;
}
std::shared_ptr<arrow::Buffer> buffer = std::move(buffer_result).ValueUnsafe();
buffer = KeepMemoryPoolAlive(std::move(buffer), pool_);
std::shared_ptr<std::atomic<uint64_t>> storage_read_bytes = storage_read_bytes_;
input_stream_->ReadAsync(
reinterpret_cast<char*>(buffer->mutable_data()), nbytes, position,
Expand Down
120 changes: 120 additions & 0 deletions src/paimon/common/utils/arrow/arrow_stream_adapter_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,11 @@
*/

#include <cstdint>
#include <cstring>
#include <functional>
#include <memory>
#include <string>
#include <utility>

#include "arrow/api.h"
#include "arrow/io/type_fwd.h"
Expand All @@ -32,6 +35,79 @@

namespace paimon::test {

namespace {

constexpr char kTestPayload[] = "data";
constexpr int64_t kTestSize = sizeof(kTestPayload) - 1;

class DeferredInputStream : public InputStream {
public:
Status Seek(int64_t, SeekOrigin) override {
return Status::OK();
}

Result<int64_t> GetPos() const override {
return 0;
}

Result<int64_t> Read(char* buffer, int64_t size) override {
if (size != kTestSize) {
return Status::Invalid("unexpected read size");
}
std::memcpy(buffer, kTestPayload, kTestSize);
return kTestSize;
}
Comment thread
mrdrivingduck marked this conversation as resolved.

Result<int64_t> Read(char* buffer, int64_t size, int64_t) override {
return Read(buffer, size);
}

void ReadAsync(char* buffer, int64_t size, int64_t,
std::function<void(Status)>&& callback) override {
buffer_ = buffer;
size_ = size;
callback_ = std::move(callback);
}

Status Complete() {
if (!callback_) {
return Status::Invalid("async request was not started");
}
if (size_ != kTestSize) {
return Status::Invalid("unexpected async read size");
}
std::memcpy(buffer_, kTestPayload, kTestSize);
auto callback = std::move(callback_);
callback(Status::OK());
return Status::OK();
}

Status Close() override {
return Status::OK();
}

Result<std::string> GetUri() const override {
return std::string("test://input");
}

Result<int64_t> Length() const override {
return kTestSize;
}

private:
char* buffer_ = nullptr;
int64_t size_ = 0;
std::function<void(Status)> callback_;
};

std::shared_ptr<ArrowInputStreamAdapter> CreateAdapter(
const std::shared_ptr<DeferredInputStream>& stream,
const std::shared_ptr<arrow::MemoryPool>& pool) {
return std::make_shared<ArrowInputStreamAdapter>(stream, kTestSize, pool);
}

} // namespace

TEST(ArrowStreamAdapterTest, TestInputAndOutputStream) {
auto test_root_dir = UniqueTestDirectory::Create();
ASSERT_TRUE(test_root_dir);
Expand Down Expand Up @@ -88,4 +164,48 @@ TEST(ArrowStreamAdapterTest, TestInputAndOutputStream) {
ASSERT_TRUE(in_stream->closed());
}

TEST(ArrowStreamAdapterTest, TestReadKeepsMemoryPoolAliveUntilBufferReleased) {
auto stream = std::make_shared<DeferredInputStream>();
std::shared_ptr<arrow::MemoryPool> pool(GetArrowPool(GetDefaultPool()));
auto adapter = CreateAdapter(stream, pool);
ASSERT_NE(adapter, nullptr);

std::shared_ptr<arrow::Buffer> buffer = adapter->Read(kTestSize).ValueOrDie();
adapter.reset();
pool.reset();

ASSERT_EQ(buffer->ToString(), kTestPayload);
buffer.reset();
}

TEST(ArrowStreamAdapterTest, TestReadAtKeepsMemoryPoolAliveUntilBufferReleased) {
auto stream = std::make_shared<DeferredInputStream>();
std::shared_ptr<arrow::MemoryPool> pool(GetArrowPool(GetDefaultPool()));
auto adapter = CreateAdapter(stream, pool);
ASSERT_NE(adapter, nullptr);

std::shared_ptr<arrow::Buffer> buffer = adapter->ReadAt(0, kTestSize).ValueOrDie();
adapter.reset();
pool.reset();

ASSERT_EQ(buffer->ToString(), kTestPayload);
buffer.reset();
}

TEST(ArrowStreamAdapterTest, TestAsyncReadKeepsMemoryPoolAliveUntilBufferReleased) {
auto stream = std::make_shared<DeferredInputStream>();
std::shared_ptr<arrow::MemoryPool> pool(GetArrowPool(GetDefaultPool()));
auto adapter = CreateAdapter(stream, pool);
ASSERT_NE(adapter, nullptr);

auto future = adapter->ReadAsync(arrow::io::default_io_context(), 0, kTestSize);
adapter.reset();
pool.reset();

ASSERT_OK(stream->Complete());
std::shared_ptr<arrow::Buffer> buffer = future.MoveResult().ValueOrDie();
ASSERT_EQ(buffer->ToString(), kTestPayload);
buffer.reset();
}

} // namespace paimon::test
Loading