mirror of
https://github.com/apache/arrow.git
synced 2026-10-28 16:47:19 +08:00
ChiLin ChiuandRaúl Cumplido
992bee2c83
### Rationale for this change See https://github.com/apache/arrow/issues/45619. ### What changes are included in this PR? Refactor using f-string instead of `string.format`. But do not use f-string for following case, `string.format` allows passing parameters, making the code more reusable. https://github.com/apache/arrow/blob/0fbf9823542233c5f32c26534c34cc97ce3f0be2/python/pyarrow/parquet/core.py#L1624-L1695 ### Are these changes tested? Via CI. ### Are there any user-facing changes? No. * GitHub Issue: #45619 Lead-authored-by: Chilin Chiou <chilin.chiou@gmail.com> Co-authored-by: Raúl Cumplido <raulcumplido@gmail.com> Signed-off-by: Raúl Cumplido <raulcumplido@gmail.com>
109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
import pyarrow as pa
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pyarrow.tests.util import rands
|
|
import memory_profiler
|
|
import gc
|
|
import io
|
|
|
|
MEGABYTE = 1 << 20
|
|
|
|
|
|
def assert_does_not_leak(f, iterations=10, check_interval=1, tolerance=5):
|
|
gc.collect()
|
|
baseline = memory_profiler.memory_usage()[0]
|
|
for i in range(iterations):
|
|
f()
|
|
if i % check_interval == 0:
|
|
gc.collect()
|
|
usage = memory_profiler.memory_usage()[0]
|
|
diff = usage - baseline
|
|
print(f"{i}: {diff}\r", end="")
|
|
if diff > tolerance:
|
|
raise Exception(f"Memory increased by {diff} megabytes after {i + 1} "
|
|
"iterations")
|
|
gc.collect()
|
|
usage = memory_profiler.memory_usage()[0]
|
|
diff = usage - baseline
|
|
print(f"\nMemory increased by {diff} megabytes after {iterations} iterations")
|
|
|
|
|
|
def test_leak1():
|
|
data = [pa.array(np.concatenate([np.random.randn(100000)] * 1000))]
|
|
table = pa.Table.from_arrays(data, ['foo'])
|
|
|
|
def func():
|
|
table.to_pandas()
|
|
assert_does_not_leak(func)
|
|
|
|
|
|
def test_leak2():
|
|
data = [pa.array(np.concatenate([np.random.randn(100000)] * 10))]
|
|
table = pa.Table.from_arrays(data, ['foo'])
|
|
|
|
def func():
|
|
df = table.to_pandas()
|
|
|
|
batch = pa.RecordBatch.from_pandas(df)
|
|
|
|
sink = io.BytesIO()
|
|
writer = pa.RecordBatchFileWriter(sink, batch.schema)
|
|
writer.write_batch(batch)
|
|
writer.close()
|
|
|
|
buf_reader = pa.BufferReader(sink.getvalue())
|
|
reader = pa.open_file(buf_reader)
|
|
reader.read_all()
|
|
|
|
assert_does_not_leak(func, iterations=50, tolerance=50)
|
|
|
|
|
|
def test_leak3():
|
|
import pyarrow.parquet as pq
|
|
|
|
df = pd.DataFrame({f'a{i}': [1, 2, 3, 4] for i in range(50)})
|
|
table = pa.Table.from_pandas(df, preserve_index=False)
|
|
|
|
writer = pq.ParquetWriter('leak_test_' + rands(5) + '.parquet',
|
|
table.schema)
|
|
|
|
def func():
|
|
writer.write_table(table, row_group_size=len(table))
|
|
|
|
# This does not "leak" per se but we do want to have this use as little
|
|
# memory as possible
|
|
assert_does_not_leak(func, iterations=500,
|
|
check_interval=50, tolerance=20)
|
|
|
|
|
|
def test_ARROW_8801():
|
|
x = pd.to_datetime(np.random.randint(0, 2**32, size=2**20, dtype=np.int64),
|
|
unit='ms', utc=True)
|
|
table = pa.table(pd.DataFrame({'x': x}))
|
|
|
|
assert_does_not_leak(lambda: table.to_pandas(split_blocks=False),
|
|
iterations=1000, check_interval=50, tolerance=1000)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_ARROW_8801()
|