@@ -1235,6 +1235,79 @@ def test_set_index(self):
1235
1235
self .assertRaises (Exception , setattr , self .mixed_frame , 'index' ,
1236
1236
idx [::2 ])
1237
1237
1238
+ def test_set_index2 (self ):
1239
+ df = DataFrame ({'A' : ['foo' , 'foo' , 'foo' , 'bar' , 'bar' ],
1240
+ 'B' : ['one' , 'two' , 'three' , 'one' , 'two' ],
1241
+ 'C' : ['a' , 'b' , 'c' , 'd' , 'e' ],
1242
+ 'D' : np .random .randn (5 ),
1243
+ 'E' : np .random .randn (5 )})
1244
+
1245
+ # new object, single-column
1246
+ result = df .set_index ('C' )
1247
+ result_nodrop = df .set_index ('C' , drop = False )
1248
+
1249
+ index = Index (df ['C' ], name = 'C' )
1250
+
1251
+ expected = df .ix [:, ['A' , 'B' , 'D' , 'E' ]]
1252
+ expected .index = index
1253
+
1254
+ expected_nodrop = df .copy ()
1255
+ expected_nodrop .index = index
1256
+
1257
+ assert_frame_equal (result , expected )
1258
+ assert_frame_equal (result_nodrop , expected_nodrop )
1259
+ self .assertEqual (result .index .name , index .name )
1260
+
1261
+ # inplace, single
1262
+ df2 = df .copy ()
1263
+ df2 .set_index ('C' , inplace = True )
1264
+ assert_frame_equal (df2 , expected )
1265
+
1266
+ df3 = df .copy ()
1267
+ df3 .set_index ('C' , drop = False , inplace = True )
1268
+ assert_frame_equal (df3 , expected_nodrop )
1269
+
1270
+ # create new object, multi-column
1271
+ result = df .set_index (['A' , 'B' ])
1272
+ result_nodrop = df .set_index (['A' , 'B' ], drop = False )
1273
+
1274
+ index = MultiIndex .from_arrays ([df ['A' ], df ['B' ]], names = ['A' , 'B' ])
1275
+
1276
+ expected = df .ix [:, ['C' , 'D' , 'E' ]]
1277
+ expected .index = index
1278
+
1279
+ expected_nodrop = df .copy ()
1280
+ expected_nodrop .index = index
1281
+
1282
+ assert_frame_equal (result , expected )
1283
+ assert_frame_equal (result_nodrop , expected_nodrop )
1284
+ self .assertEqual (result .index .names , index .names )
1285
+
1286
+ # inplace
1287
+ df2 = df .copy ()
1288
+ df2 .set_index (['A' , 'B' ], inplace = True )
1289
+ assert_frame_equal (df2 , expected )
1290
+
1291
+ df3 = df .copy ()
1292
+ df3 .set_index (['A' , 'B' ], drop = False , inplace = True )
1293
+ assert_frame_equal (df3 , expected_nodrop )
1294
+
1295
+ # corner case
1296
+ self .assertRaises (Exception , df .set_index , 'A' , verify_integrity = True )
1297
+
1298
+ def test_set_index_pass_arrays (self ):
1299
+ df = DataFrame ({'A' : ['foo' , 'bar' , 'foo' , 'bar' ,
1300
+ 'foo' , 'bar' , 'foo' , 'foo' ],
1301
+ 'B' : ['one' , 'one' , 'two' , 'three' ,
1302
+ 'two' , 'two' , 'one' , 'three' ],
1303
+ 'C' : np .random .randn (8 ),
1304
+ 'D' : np .random .randn (8 )})
1305
+
1306
+ # multiple columns
1307
+ result = df .set_index (['A' , df ['B' ].values ], drop = False )
1308
+ expected = df .set_index (['A' , 'B' ], drop = False )
1309
+ assert_frame_equal (result , expected )
1310
+
1238
1311
def test_set_columns (self ):
1239
1312
cols = Index (np .arange (len (self .mixed_frame .columns )))
1240
1313
self .mixed_frame .columns = cols
@@ -3502,66 +3575,6 @@ def test_reindex_fill_value(self):
3502
3575
expected = df .reindex (range (15 )).fillna (0 )
3503
3576
assert_frame_equal (result , expected )
3504
3577
3505
- def test_set_index2 (self ):
3506
- df = DataFrame ({'A' : ['foo' , 'foo' , 'foo' , 'bar' , 'bar' ],
3507
- 'B' : ['one' , 'two' , 'three' , 'one' , 'two' ],
3508
- 'C' : ['a' , 'b' , 'c' , 'd' , 'e' ],
3509
- 'D' : np .random .randn (5 ),
3510
- 'E' : np .random .randn (5 )})
3511
-
3512
- # new object, single-column
3513
- result = df .set_index ('C' )
3514
- result_nodrop = df .set_index ('C' , drop = False )
3515
-
3516
- index = Index (df ['C' ], name = 'C' )
3517
-
3518
- expected = df .ix [:, ['A' , 'B' , 'D' , 'E' ]]
3519
- expected .index = index
3520
-
3521
- expected_nodrop = df .copy ()
3522
- expected_nodrop .index = index
3523
-
3524
- assert_frame_equal (result , expected )
3525
- assert_frame_equal (result_nodrop , expected_nodrop )
3526
- self .assertEqual (result .index .name , index .name )
3527
-
3528
- # inplace, single
3529
- df2 = df .copy ()
3530
- df2 .set_index ('C' , inplace = True )
3531
- assert_frame_equal (df2 , expected )
3532
-
3533
- df3 = df .copy ()
3534
- df3 .set_index ('C' , drop = False , inplace = True )
3535
- assert_frame_equal (df3 , expected_nodrop )
3536
-
3537
- # create new object, multi-column
3538
- result = df .set_index (['A' , 'B' ])
3539
- result_nodrop = df .set_index (['A' , 'B' ], drop = False )
3540
-
3541
- index = MultiIndex .from_arrays ([df ['A' ], df ['B' ]], names = ['A' , 'B' ])
3542
-
3543
- expected = df .ix [:, ['C' , 'D' , 'E' ]]
3544
- expected .index = index
3545
-
3546
- expected_nodrop = df .copy ()
3547
- expected_nodrop .index = index
3548
-
3549
- assert_frame_equal (result , expected )
3550
- assert_frame_equal (result_nodrop , expected_nodrop )
3551
- self .assertEqual (result .index .names , index .names )
3552
-
3553
- # inplace
3554
- df2 = df .copy ()
3555
- df2 .set_index (['A' , 'B' ], inplace = True )
3556
- assert_frame_equal (df2 , expected )
3557
-
3558
- df3 = df .copy ()
3559
- df3 .set_index (['A' , 'B' ], drop = False , inplace = True )
3560
- assert_frame_equal (df3 , expected_nodrop )
3561
-
3562
- # corner case
3563
- self .assertRaises (Exception , df .set_index , 'A' )
3564
-
3565
3578
def test_align (self ):
3566
3579
af , bf = self .frame .align (self .frame )
3567
3580
self .assert_ (af ._data is not self .frame ._data )
0 commit comments