Skip to content
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

Jump to bottom

[c++] Fix memory leak in Arrow table implementation #6314

Merged
merged 3 commits into from
Feb 22, 2024
Merged

[c++] Fix memory leak in Arrow table implementation #6314

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4ef35a0
[c++] Fix memory leak in Arrow table implementation
borchero Feb 14, 2024
f7ec643
Merge branch 'master' into fix-arrow-leak
jameslamb Feb 20, 2024
bfb1b67
Merge branch 'master' into fix-arrow-leak
borchero Feb 22, 2024
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
21 changes: 20 additions & 1 deletion include/LightGBM/arrow.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,9 @@ class ArrowChunkedArray {
*/
class ArrowTable {
std::vector<ArrowChunkedArray> columns_;
const int64_t n_chunks_;
const ArrowArray* chunks_ptr_;
const ArrowSchema* schema_ptr_;

public:
/**
Expand All @@ -216,7 +219,8 @@ class ArrowTable {
* @param chunks A C-style array containing the chunks.
* @param schema The schema for all chunks.
*/
inline ArrowTable(int64_t n_chunks, const ArrowArray* chunks, const ArrowSchema* schema) {
inline ArrowTable(int64_t n_chunks, const ArrowArray *chunks, const ArrowSchema *schema)
: n_chunks_(n_chunks), chunks_ptr_(chunks), schema_ptr_(schema) {
columns_.reserve(schema->n_children);
for (int64_t j = 0; j < schema->n_children; ++j) {
std::vector<const ArrowArray*> children_chunks;
Expand All @@ -229,6 +233,21 @@ class ArrowTable {
}
}

~ArrowTable() {
// As consumer of the Arrow array, the Arrow table must release all Arrow arrays it receives
// as well as the schema. As per the specification, children arrays are released by the
// producer. See: https://arrow.apache.org/docs/format/CDataInterface.html#release-callback-semantics-for-consumers
Copy link
Collaborator

Choose a reason for hiding this comment

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

thanks for including this link, super helpful!

for (int64_t i = 0; i < n_chunks_; ++i) {
auto chunk = &chunks_ptr_[i];
if (chunk->release) {
chunk->release(const_cast<ArrowArray*>(chunk));
}
}
if (schema_ptr_->release) {
schema_ptr_->release(const_cast<ArrowSchema*>(schema_ptr_));
}
}

/**
* @brief Get the number of rows in the table.
*
Expand Down
Loading