Skip to content

Commit 46cf332

Browse files
committed
ARROW-8801: [Python] Fix memory leak when converting datetime64-with-tz data to pandas
An array object was failing to be decref'd on the DatetimeTZ conversion path. The code is slightly complicated by the different calling reference ownership semantics of the Array/ChunkedArray conversion path (which expects to own the created array when it calls `GetSeriesResult` while the `GetResultBlock` code retains its array in a `OwnedRefNoGIL`). This was the simplest thing that fixed the memory leak for me. There is leak checking code that can be used to verify this in python/scripts/test_leak.py (just run the script). Closes apache#7522 from wesm/ARROW-8801 Authored-by: Wes McKinney <wesm@apache.org> Signed-off-by: Wes McKinney <wesm@apache.org>
1 parent 7cebe08 commit 46cf332

2 files changed

Lines changed: 34 additions & 6 deletions

File tree

‎cpp/src/arrow/python/arrow_to_pandas.cc‎

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -403,12 +403,18 @@ class PandasWriter {
403403
return Status::OK();
404404
}
405405

406-
virtual Status GetSeriesResult(PyObject** out) { return GetBlock1D(out); }
406+
// Caller steals the reference to this object
407+
virtual Status GetSeriesResult(PyObject** out) {
408+
RETURN_NOT_OK(MakeBlock1D());
409+
// Caller owns the object now
410+
*out = block_arr_.detach();
411+
return Status::OK();
412+
}
407413

408414
protected:
409415
virtual Status AddResultMetadata(PyObject* result) { return Status::OK(); }
410416

411-
Status GetBlock1D(PyObject** out) {
417+
Status MakeBlock1D() {
412418
// For Series or for certain DataFrame block types, we need to shape to a
413419
// 1D array when there is only one column
414420
PyAcquireGIL lock;
@@ -420,9 +426,17 @@ class PandasWriter {
420426
dims.ptr = new_dims;
421427
dims.len = 1;
422428

423-
*out = PyArray_Newshape(reinterpret_cast<PyArrayObject*>(block_arr_.obj()), &dims,
424-
NPY_ANYORDER);
429+
PyObject* reshaped = PyArray_Newshape(
430+
reinterpret_cast<PyArrayObject*>(block_arr_.obj()), &dims, NPY_ANYORDER);
425431
RETURN_IF_PYERROR();
432+
433+
// ARROW-8801: Here a PyArrayObject is created that is not being managed by
434+
// any OwnedRef object. This object is then put in the resulting object
435+
// with PyDict_SetItemString, which increments the reference count, so a
436+
// memory leak ensues. There are several ways to fix the memory leak but a
437+
// simple one is to put the reshaped 1D block array in this OwnedRefNoGIL
438+
// so it will be correctly decref'd when this class is destructed.
439+
block_arr_.reset(reshaped);
426440
return Status::OK();
427441
}
428442

@@ -1269,13 +1283,18 @@ class DatetimeTZWriter : public DatetimeNanoWriter {
12691283
: DatetimeNanoWriter(options, num_rows, 1), timezone_(timezone) {}
12701284

12711285
protected:
1272-
Status GetResultBlock(PyObject** out) override { return GetBlock1D(out); }
1286+
Status GetResultBlock(PyObject** out) override {
1287+
RETURN_NOT_OK(MakeBlock1D());
1288+
*out = block_arr_.obj();
1289+
return Status::OK();
1290+
}
12731291

12741292
Status AddResultMetadata(PyObject* result) override {
12751293
PyObject* py_tz = PyUnicode_FromStringAndSize(
12761294
timezone_.c_str(), static_cast<Py_ssize_t>(timezone_.size()));
12771295
RETURN_IF_PYERROR();
12781296
PyDict_SetItemString(result, "timezone", py_tz);
1297+
Py_DECREF(py_tz);
12791298
return Status::OK();
12801299
}
12811300

‎python/scripts/test_leak.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,5 +97,14 @@ def func():
9797
check_interval=50, tolerance=20)
9898

9999

100+
def test_ARROW_8801():
101+
x = pd.to_datetime(np.random.randint(0, 2**32, size=2**20),
102+
unit='ms', utc=True)
103+
table = pa.table(pd.DataFrame({'x': x}))
104+
105+
assert_does_not_leak(lambda: table.to_pandas(split_blocks=False),
106+
iterations=1000, check_interval=50, tolerance=1000)
107+
108+
100109
if __name__ == '__main__':
101-
test_leak3()
110+
test_ARROW_8801()

0 commit comments

Comments
 (0)