chunyang opened a new pull request, #40412:
URL: https://github.com/apache/arrow/pull/40412
<!--
Thanks for opening a pull request!
If this is your first pull request you can find detailed information on how
to contribute here:
* [New Contributor's
Guide](https://arrow.apache.org/docs/dev/developers/guide/step_by_step/pr_lifecycle.html#reviews-and-merge-of-the-pull-request)
* [Contributing
Overview](https://arrow.apache.org/docs/dev/developers/overview.html)
If this is not a [minor
PR](https://github.com/apache/arrow/blob/main/CONTRIBUTING.md#Minor-Fixes).
Could you open an issue for this pull request on GitHub?
https://github.com/apache/arrow/issues/new/choose
Opening GitHub issues ahead of time contributes to the
[Openness](http://theapacheway.com/open/#:~:text=Openness%20allows%20new%20users%20the,must%20happen%20in%20the%20open.)
of the Apache Arrow project.
Then could you also rename the pull request title in the following format?
GH-${GITHUB_ISSUE_ID}: [${COMPONENT}] ${SUMMARY}
or
MINOR: [${COMPONENT}] ${SUMMARY}
In the case of PARQUET issues on JIRA the title also supports:
PARQUET-${JIRA_ISSUE_ID}: [${COMPONENT}] ${SUMMARY}
-->
### Rationale for this change
When creating Arrow arrays using `pa.array` from lists of dicts, memory
usage is observed to increase over time even 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.
<!--
Why are you proposing this change? If this is already explained clearly in
the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your
changes and offer better suggestions for fixes.
-->
### 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`.
<!--
There is no need to duplicate the description in the issue here but it is
sometimes worth providing a summary of the individual changes in this PR.
-->
### 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
https://github.com/apache/arrow/issues/37989#issue-1924129600, 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.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are
they covered by existing tests)?
-->
### Are there any user-facing changes?
None
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please uncomment the line
below and explain which changes are breaking.
-->
<!-- **This PR includes breaking changes to public APIs.** -->
<!--
Please uncomment the line below (and provide explanation) if the changes fix
either (a) a security vulnerability, (b) a bug that caused incorrect or invalid
data to be produced, or (c) a bug that causes a crash (even when the API
contract is upheld). We use this to highlight fixes to issues that may affect
users without their knowledge. For this reason, fixing bugs that cause errors
don't count, since those are usually obvious.
-->
<!-- **This PR contains a "Critical Fix".** -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
