Skip to content

Commit aeda02e

Browse files
sfc-gh-stakedaankit-bhatnagar167
authored andcommitted
SNOW-98856: Fix memory leak in python connector panda df fetch API
1 parent 7e45a86 commit aeda02e

File tree

2 files changed

+13
-2
lines changed

2 files changed

+13
-2
lines changed

cpp/ArrowIterator/CArrowTableIterator.cpp

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ void CArrowTableIterator::reconstructRecordBatches()
128128
}
129129

130130
CArrowTableIterator::CArrowTableIterator(PyObject* context, PyObject* batches)
131-
: CArrowIterator(batches), m_context(context)
131+
: CArrowIterator(batches), m_context(context), m_pyTableObjRef(nullptr)
132132
{
133133
PyObject* tz = PyObject_GetAttrString(m_context, "_timezone");
134134
PyArg_Parse(tz, "s", &m_timezone);
@@ -138,7 +138,15 @@ CArrowTableIterator::CArrowTableIterator(PyObject* context, PyObject* batches)
138138
PyObject* CArrowTableIterator::next()
139139
{
140140
bool firstDone = this->convertRecordBatchesToTable();
141-
return (firstDone && m_cTable) ? arrow::py::wrap_table(m_cTable) : Py_None;
141+
if (firstDone && m_cTable)
142+
{
143+
m_pyTableObjRef.reset(arrow::py::wrap_table(m_cTable));
144+
return m_pyTableObjRef.get();
145+
}
146+
else
147+
{
148+
return Py_None;
149+
}
142150
}
143151

144152
arrow::Status CArrowTableIterator::replaceColumn(

cpp/ArrowIterator/CArrowTableIterator.hpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ class CArrowTableIterator : public CArrowIterator
4141
/** arrow format convert context for the current session */
4242
PyObject* m_context;
4343

44+
/** reference to PyObject */
45+
py::UniqueRef m_pyTableObjRef;
46+
4447
/**
4548
* arrow memory buffer to allocate type converted arrays for fetching pandas from arrow
4649
*/

0 commit comments

Comments
 (0)