@@ -250,3 +250,33 @@ def test_set_index(using_copy_on_write):
250
250
df2 .iloc [0 , 1 ] = 0
251
251
assert not np .shares_memory (get_array (df2 , "c" ), get_array (df , "c" ))
252
252
tm .assert_frame_equal (df , df_orig )
253
+
254
+
255
+ @pytest .mark .parametrize (
256
+ "method" ,
257
+ [
258
+ lambda df : df .head (),
259
+ lambda df : df .head (2 ),
260
+ lambda df : df .tail (),
261
+ lambda df : df .tail (3 ),
262
+ ],
263
+ )
264
+ def test_head_tail (method , using_copy_on_write ):
265
+ df = DataFrame ({"a" : [1 , 2 , 3 ], "b" : [0.1 , 0.2 , 0.3 ]})
266
+ df_orig = df .copy ()
267
+ df2 = method (df )
268
+ df2 ._mgr ._verify_integrity ()
269
+
270
+ if using_copy_on_write :
271
+ assert np .shares_memory (get_array (df2 , "a" ), get_array (df , "a" ))
272
+ assert np .shares_memory (get_array (df2 , "b" ), get_array (df , "b" ))
273
+
274
+ # modify df2 to trigger CoW for that block
275
+ df2 .iloc [0 , 0 ] = 0
276
+ assert np .shares_memory (get_array (df2 , "b" ), get_array (df , "b" ))
277
+ if using_copy_on_write :
278
+ assert not np .shares_memory (get_array (df2 , "a" ), get_array (df , "a" ))
279
+ else :
280
+ # without CoW enabled, head and tail return views. Mutating df2 also mutates df.
281
+ df2 .iloc [0 , 0 ] = 1
282
+ tm .assert_frame_equal (df , df_orig )
0 commit comments