GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts #40412

chunyang · 2024-03-07T23:22:13Z

Rationale for this change

When creating Arrow arrays using pa.array from lists of dicts, memory usage is observed to increase over time despite the created arrays going out of scope. The issue appears to only happen for lists of dicts, as opposed to lists of numpy arrays or other types.

What changes are included in this PR?

This PR makes two changes to python_to_arrow.cc, to ensure that new references created by PyDict_Items and PySequence_GetItem are properly reference counted via OwnedRef.

Are these changes tested?

The change was tested against the following reproduction script:

"""Repro memory increase observed when creating pyarrow arrays."""

# System imports
import logging

# Third-party imports
import numpy as np
import psutil
import pyarrow as pa

LIST_LENGTH = 5 * (2**20)
LOGGER = logging.getLogger(__name__)


def initialize_logging() -> None:
    logging.basicConfig(
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        level=logging.INFO,
    )


def get_rss_in_mib() -> float:
    """Return the Resident Set Size of the current process in MiB."""
    return psutil.Process().memory_info().rss / 1024 / 1024


def main() -> None:
    initialize_logging()

    for idx in range(100):
        data = np.random.randint(256, size=(LIST_LENGTH,), dtype=np.uint8)
        # data = "a" * LIST_LENGTH
        pa.array([{"data": data}])
        if (idx + 1) % 10 == 0:
            LOGGER.info(
                "%d dict arrays created, RSS: %.2f MiB", idx + 1, get_rss_in_mib()
            )

    LOGGER.info("---------")

    for idx in range(100):
        pa.array(
            [
                np.random.randint(256, size=(LIST_LENGTH,), dtype=np.uint8).tobytes(),
            ]
        )
        if (idx + 1) % 10 == 0:
            LOGGER.info(
                "%d non-dict arrays created, RSS: %.2f MiB", idx + 1, get_rss_in_mib()
            )


if __name__ == "__main__":
    main()

Prior to this change, the reproduction script produces the following output:

2024-03-07 23:14:17,560 - __main__ - INFO - 10 dict arrays created, RSS: 121.05 MiB
2024-03-07 23:14:17,698 - __main__ - INFO - 20 dict arrays created, RSS: 171.07 MiB
2024-03-07 23:14:17,835 - __main__ - INFO - 30 dict arrays created, RSS: 221.09 MiB
2024-03-07 23:14:17,971 - __main__ - INFO - 40 dict arrays created, RSS: 271.11 MiB
2024-03-07 23:14:18,109 - __main__ - INFO - 50 dict arrays created, RSS: 320.86 MiB
2024-03-07 23:14:18,245 - __main__ - INFO - 60 dict arrays created, RSS: 371.65 MiB
2024-03-07 23:14:18,380 - __main__ - INFO - 70 dict arrays created, RSS: 422.18 MiB
2024-03-07 23:14:18,516 - __main__ - INFO - 80 dict arrays created, RSS: 472.20 MiB
2024-03-07 23:14:18,650 - __main__ - INFO - 90 dict arrays created, RSS: 522.21 MiB
2024-03-07 23:14:18,788 - __main__ - INFO - 100 dict arrays created, RSS: 572.23 MiB
2024-03-07 23:14:18,789 - __main__ - INFO - ---------
2024-03-07 23:14:19,001 - __main__ - INFO - 10 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:19,211 - __main__ - INFO - 20 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:19,417 - __main__ - INFO - 30 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:19,623 - __main__ - INFO - 40 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:19,832 - __main__ - INFO - 50 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:20,047 - __main__ - INFO - 60 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:20,253 - __main__ - INFO - 70 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:20,499 - __main__ - INFO - 80 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:20,725 - __main__ - INFO - 90 non-dict arrays created, RSS: 567.61 MiB
2024-03-07 23:14:20,950 - __main__ - INFO - 100 non-dict arrays created, RSS: 567.61 MiB

After this change, the output changes to the following. Notice that the Resident Set Size (RSS) no longer increases as more Arrow arrays are created from list of dict.

2024-03-07 23:14:47,246 - __main__ - INFO - 10 dict arrays created, RSS: 81.73 MiB
2024-03-07 23:14:47,353 - __main__ - INFO - 20 dict arrays created, RSS: 76.53 MiB
2024-03-07 23:14:47,445 - __main__ - INFO - 30 dict arrays created, RSS: 82.20 MiB
2024-03-07 23:14:47,537 - __main__ - INFO - 40 dict arrays created, RSS: 86.59 MiB
2024-03-07 23:14:47,634 - __main__ - INFO - 50 dict arrays created, RSS: 80.28 MiB
2024-03-07 23:14:47,734 - __main__ - INFO - 60 dict arrays created, RSS: 85.44 MiB
2024-03-07 23:14:47,827 - __main__ - INFO - 70 dict arrays created, RSS: 85.44 MiB
2024-03-07 23:14:47,921 - __main__ - INFO - 80 dict arrays created, RSS: 85.44 MiB
2024-03-07 23:14:48,024 - __main__ - INFO - 90 dict arrays created, RSS: 82.94 MiB
2024-03-07 23:14:48,132 - __main__ - INFO - 100 dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,132 - __main__ - INFO - ---------
2024-03-07 23:14:48,229 - __main__ - INFO - 10 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,324 - __main__ - INFO - 20 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,420 - __main__ - INFO - 30 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,516 - __main__ - INFO - 40 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,613 - __main__ - INFO - 50 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,710 - __main__ - INFO - 60 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,806 - __main__ - INFO - 70 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:48,905 - __main__ - INFO - 80 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:49,009 - __main__ - INFO - 90 non-dict arrays created, RSS: 87.84 MiB
2024-03-07 23:14:49,108 - __main__ - INFO - 100 non-dict arrays created, RSS: 87.84 MiB

When this change is tested against the reproduction script provided in #37989 (comment), the reported memory increase is no longer observed.

I have not added a unit test, but it may be possible to add one similar to the reproduction scripts used above, provided there's an accurate way to capture process memory usage on all the platforms that Arrow supports, and provided memory usage is not affected by concurrently running tests. If this code could be tested under valgrind, that may be an even better way to go.

Are there any user-facing changes?

GitHub Issue: [Python] pyarrow.Table.from_pandas() causing memory leak #37989

github-actions · 2024-03-07T23:22:38Z

⚠️ GitHub issue #37989 has been automatically assigned in GitHub to PR creator.

github-actions · 2024-03-07T23:24:18Z

⚠️ GitHub issue #37989 has been automatically assigned in GitHub to PR creator.

Ashokcs94

Looks fine for me

kou

+1

pitrou · 2024-03-08T09:47:52Z

python/pyarrow/src/arrow/python/python_to_arrow.cc

@@ -1114,6 +1115,7 @@ class PyStructConverter : public StructConverter<PyConverter, PyConverterTrait>

  Result<std::pair<PyObject*, PyObject*>> GetKeyValuePair(PyObject* seq, int index) {
    PyObject* pair = PySequence_GetItem(seq, index);
+    OwnedRef pair_ref(pair);


Should be moved just below RETURN_IF_PYERROR.

Thanks for the review! I made the suggested move.

pitrou

Thanks a lot for diagnosing this. This is a potentially large leak...

jorisvandenbossche

Thanks a lot!

conbench-apache-arrow · 2024-03-17T02:05:52Z

After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit 03c771a.

There were no benchmark performance regressions. 🎉

The full Conbench report has more details. It also includes information about 3 possible false positives for unstable benchmarks that are known to sometimes produce them.

galipremsagar · 2024-04-15T17:55:36Z

@pitrou | @jorisvandenbossche Due to the seriousness of the leak, is it possible to back-port this fix for 14.x and 15.x versions of pyarrow?

I opened two backport PRs if you want to go ahead:

…ay from Python list of dicts (apache#40412) ### Rationale for this change When creating Arrow arrays using `pa.array` from lists of dicts, memory usage is observed to increase over time despite the created arrays going out of scope. The issue appears to only happen for lists of dicts, as opposed to lists of numpy arrays or other types. ### What changes are included in this PR? This PR makes two changes to _python_to_arrow.cc_, to ensure that new references created by [`PyDict_Items`](https://docs.python.org/3/c-api/dict.html#c.PyDict_Items) and [`PySequence_GetItem`](https://docs.python.org/3/c-api/sequence.html#c.PySequence_GetItem) are properly reference counted via `OwnedRef`. ### Are these changes tested? The change was tested against the following reproduction script: ```python """Repro memory increase observed when creating pyarrow arrays.""" # System imports import logging # Third-party imports import numpy as np import psutil import pyarrow as pa LIST_LENGTH = 5 * (2**20) LOGGER = logging.getLogger(__name__) def initialize_logging() -> None: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO, ) def get_rss_in_mib() -> float: """Return the Resident Set Size of the current process in MiB.""" return psutil.Process().memory_info().rss / 1024 / 1024 def main() -> None: initialize_logging() for idx in range(100): data = np.random.randint(256, size=(LIST_LENGTH,), dtype=np.uint8) # data = "a" * LIST_LENGTH pa.array([{"data": data}]) if (idx + 1) % 10 == 0: LOGGER.info( "%d dict arrays created, RSS: %.2f MiB", idx + 1, get_rss_in_mib() ) LOGGER.info("---------") for idx in range(100): pa.array( [ np.random.randint(256, size=(LIST_LENGTH,), dtype=np.uint8).tobytes(), ] ) if (idx + 1) % 10 == 0: LOGGER.info( "%d non-dict arrays created, RSS: %.2f MiB", idx + 1, get_rss_in_mib() ) if __name__ == "__main__": main() ``` Prior to this change, the reproduction script produces the following output: ``` 2024-03-07 23:14:17,560 - __main__ - INFO - 10 dict arrays created, RSS: 121.05 MiB 2024-03-07 23:14:17,698 - __main__ - INFO - 20 dict arrays created, RSS: 171.07 MiB 2024-03-07 23:14:17,835 - __main__ - INFO - 30 dict arrays created, RSS: 221.09 MiB 2024-03-07 23:14:17,971 - __main__ - INFO - 40 dict arrays created, RSS: 271.11 MiB 2024-03-07 23:14:18,109 - __main__ - INFO - 50 dict arrays created, RSS: 320.86 MiB 2024-03-07 23:14:18,245 - __main__ - INFO - 60 dict arrays created, RSS: 371.65 MiB 2024-03-07 23:14:18,380 - __main__ - INFO - 70 dict arrays created, RSS: 422.18 MiB 2024-03-07 23:14:18,516 - __main__ - INFO - 80 dict arrays created, RSS: 472.20 MiB 2024-03-07 23:14:18,650 - __main__ - INFO - 90 dict arrays created, RSS: 522.21 MiB 2024-03-07 23:14:18,788 - __main__ - INFO - 100 dict arrays created, RSS: 572.23 MiB 2024-03-07 23:14:18,789 - __main__ - INFO - --------- 2024-03-07 23:14:19,001 - __main__ - INFO - 10 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:19,211 - __main__ - INFO - 20 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:19,417 - __main__ - INFO - 30 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:19,623 - __main__ - INFO - 40 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:19,832 - __main__ - INFO - 50 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:20,047 - __main__ - INFO - 60 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:20,253 - __main__ - INFO - 70 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:20,499 - __main__ - INFO - 80 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:20,725 - __main__ - INFO - 90 non-dict arrays created, RSS: 567.61 MiB 2024-03-07 23:14:20,950 - __main__ - INFO - 100 non-dict arrays created, RSS: 567.61 MiB ``` After this change, the output changes to the following. Notice that the Resident Set Size (RSS) no longer increases as more Arrow arrays are created from list of dict. ``` 2024-03-07 23:14:47,246 - __main__ - INFO - 10 dict arrays created, RSS: 81.73 MiB 2024-03-07 23:14:47,353 - __main__ - INFO - 20 dict arrays created, RSS: 76.53 MiB 2024-03-07 23:14:47,445 - __main__ - INFO - 30 dict arrays created, RSS: 82.20 MiB 2024-03-07 23:14:47,537 - __main__ - INFO - 40 dict arrays created, RSS: 86.59 MiB 2024-03-07 23:14:47,634 - __main__ - INFO - 50 dict arrays created, RSS: 80.28 MiB 2024-03-07 23:14:47,734 - __main__ - INFO - 60 dict arrays created, RSS: 85.44 MiB 2024-03-07 23:14:47,827 - __main__ - INFO - 70 dict arrays created, RSS: 85.44 MiB 2024-03-07 23:14:47,921 - __main__ - INFO - 80 dict arrays created, RSS: 85.44 MiB 2024-03-07 23:14:48,024 - __main__ - INFO - 90 dict arrays created, RSS: 82.94 MiB 2024-03-07 23:14:48,132 - __main__ - INFO - 100 dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,132 - __main__ - INFO - --------- 2024-03-07 23:14:48,229 - __main__ - INFO - 10 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,324 - __main__ - INFO - 20 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,420 - __main__ - INFO - 30 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,516 - __main__ - INFO - 40 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,613 - __main__ - INFO - 50 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,710 - __main__ - INFO - 60 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,806 - __main__ - INFO - 70 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:48,905 - __main__ - INFO - 80 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:49,009 - __main__ - INFO - 90 non-dict arrays created, RSS: 87.84 MiB 2024-03-07 23:14:49,108 - __main__ - INFO - 100 non-dict arrays created, RSS: 87.84 MiB ``` When this change is tested against the reproduction script provided in apache#37989 (comment), the reported memory increase is no longer observed. I have not added a unit test, but it may be possible to add one similar to the reproduction scripts used above, provided there's an accurate way to capture process memory usage on all the platforms that Arrow supports, and provided memory usage is not affected by concurrently running tests. If this code could be tested under valgrind, that may be an even better way to go. ### Are there any user-facing changes? * GitHub Issue: apache#37989 Authored-by: Chuck Yang <chuck.yang@getcruise.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>

jorisvandenbossche · 2024-04-16T07:31:14Z

@galipremsagar we are currently in the process of releasing pyarrow 16.0, and I think it is quite unlikely that we will still do a patch release for 14.0 and 15.0 (given the work involved to do a release at the moment)

jakirkham · 2024-04-16T08:31:43Z

In that case, maybe we can just patch the conda-forge packages

Thought this might help more users with this issue (like wheel consumers), but understand if that is not feasible

Plug a couple of reference leaks

0fbedc1

github-actions bot added Component: Python awaiting review Awaiting review labels Mar 7, 2024

chunyang mentioned this pull request Mar 7, 2024

[Python] pyarrow.Table.from_pandas() causing memory leak #37989

Closed

chunyang changed the title ~~GH-37989: [C++] Plug reference leaks when creating Arrow array from Python list of dicts~~ GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts Mar 7, 2024

Ashokcs94 approved these changes Mar 8, 2024

View reviewed changes

kou approved these changes Mar 8, 2024

View reviewed changes

github-actions bot added awaiting committer review Awaiting committer review and removed awaiting review Awaiting review labels Mar 8, 2024

kou requested review from jorisvandenbossche and pitrou March 8, 2024 05:20

github-actions bot added awaiting merge Awaiting merge and removed awaiting committer review Awaiting committer review labels Mar 8, 2024

pitrou reviewed Mar 8, 2024

View reviewed changes

pitrou requested changes Mar 8, 2024

View reviewed changes

chunyang added 2 commits March 8, 2024 10:13

Move OwnedRef creation after RETURN_IF_PYERROR

8dcfe4d

Reword inline comment

84b635c

jorisvandenbossche approved these changes Mar 15, 2024

View reviewed changes

jorisvandenbossche merged commit 03c771a into apache:main Mar 15, 2024
12 of 13 checks passed

jorisvandenbossche removed the awaiting merge Awaiting merge label Mar 15, 2024

chunyang deleted the cyang/struct-mem-leak branch March 15, 2024 17:44

wence- mentioned this pull request Apr 15, 2024

[Python] Memory leak for repeated StructArray creation #41172

Closed

This was referenced Apr 15, 2024

Backport #40412 to pyarrow-15.x #41221

Closed

Backport #40412 to pyarrow-14.x #41222

Closed

This was referenced Apr 18, 2024

Backport #40412 to 15.0.2 conda-forge/arrow-cpp-feedstock#1368

Merged

Backport #40412 to 14.0.2 conda-forge/arrow-cpp-feedstock#1369

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts #40412

GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts #40412

chunyang commented Mar 7, 2024 •

edited

Loading

github-actions bot commented Mar 7, 2024

github-actions bot commented Mar 7, 2024

Ashokcs94 left a comment

kou left a comment

pitrou Mar 8, 2024

chunyang Mar 8, 2024

pitrou left a comment

jorisvandenbossche left a comment

conbench-apache-arrow bot commented Mar 17, 2024

galipremsagar commented Apr 15, 2024 •

edited

Loading

jorisvandenbossche commented Apr 16, 2024

jakirkham commented Apr 16, 2024

GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts #40412

GH-37989: [Python] Plug reference leaks when creating Arrow array from Python list of dicts #40412

Conversation

chunyang commented Mar 7, 2024 • edited Loading

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

github-actions bot commented Mar 7, 2024

github-actions bot commented Mar 7, 2024

Ashokcs94 left a comment

Choose a reason for hiding this comment

kou left a comment

Choose a reason for hiding this comment

pitrou Mar 8, 2024

Choose a reason for hiding this comment

chunyang Mar 8, 2024

Choose a reason for hiding this comment

pitrou left a comment

Choose a reason for hiding this comment

jorisvandenbossche left a comment

Choose a reason for hiding this comment

conbench-apache-arrow bot commented Mar 17, 2024

galipremsagar commented Apr 15, 2024 • edited Loading

jorisvandenbossche commented Apr 16, 2024

jakirkham commented Apr 16, 2024

chunyang commented Mar 7, 2024 •

edited

Loading

galipremsagar commented Apr 15, 2024 •

edited

Loading