import numpy as np
|
import pytest
|
|
from pandas import DataFrame
|
import pandas._testing as tm
|
from pandas.tests.copy_view.util import get_array
|
|
|
def test_assigning_to_same_variable_removes_references(using_copy_on_write):
|
df = DataFrame({"a": [1, 2, 3]})
|
df = df.reset_index()
|
if using_copy_on_write:
|
assert df._mgr._has_no_reference(1)
|
arr = get_array(df, "a")
|
df.iloc[0, 1] = 100 # Write into a
|
|
assert np.shares_memory(arr, get_array(df, "a"))
|
|
|
def test_setitem_dont_track_unnecessary_references(using_copy_on_write):
|
df = DataFrame({"a": [1, 2, 3], "b": 1, "c": 1})
|
|
df["b"] = 100
|
arr = get_array(df, "a")
|
# We split the block in setitem, if we are not careful the new blocks will
|
# reference each other triggering a copy
|
df.iloc[0, 0] = 100
|
assert np.shares_memory(arr, get_array(df, "a"))
|
|
|
def test_setitem_with_view_copies(using_copy_on_write):
|
df = DataFrame({"a": [1, 2, 3], "b": 1, "c": 1})
|
view = df[:]
|
expected = df.copy()
|
|
df["b"] = 100
|
arr = get_array(df, "a")
|
df.iloc[0, 0] = 100 # Check that we correctly track reference
|
if using_copy_on_write:
|
assert not np.shares_memory(arr, get_array(df, "a"))
|
tm.assert_frame_equal(view, expected)
|
|
|
def test_setitem_with_view_invalidated_does_not_copy(using_copy_on_write, request):
|
df = DataFrame({"a": [1, 2, 3], "b": 1, "c": 1})
|
view = df[:]
|
|
df["b"] = 100
|
arr = get_array(df, "a")
|
view = None # noqa
|
df.iloc[0, 0] = 100
|
if using_copy_on_write:
|
# Setitem split the block. Since the old block shared data with view
|
# all the new blocks are referencing view and each other. When view
|
# goes out of scope, they don't share data with any other block,
|
# so we should not trigger a copy
|
mark = pytest.mark.xfail(
|
reason="blk.delete does not track references correctly"
|
)
|
request.node.add_marker(mark)
|
assert np.shares_memory(arr, get_array(df, "a"))
|
|
|
def test_out_of_scope(using_copy_on_write):
|
def func():
|
df = DataFrame({"a": [1, 2], "b": 1.5, "c": 1})
|
# create some subset
|
result = df[["a", "b"]]
|
return result
|
|
result = func()
|
if using_copy_on_write:
|
assert not result._mgr.blocks[0].refs.has_reference()
|
assert not result._mgr.blocks[1].refs.has_reference()
|
|
|
def test_delete(using_copy_on_write):
|
df = DataFrame(np.random.randn(4, 3), columns=["a", "b", "c"])
|
del df["b"]
|
if using_copy_on_write:
|
# TODO: This should not have references, delete makes a shallow copy
|
# but keeps the blocks alive
|
assert df._mgr.blocks[0].refs.has_reference()
|
assert df._mgr.blocks[1].refs.has_reference()
|
|
df = df[["a"]]
|
if using_copy_on_write:
|
assert not df._mgr.blocks[0].refs.has_reference()
|
