DEV Community

wrighter
wrighter

Posted on • Originally published at wrighters.io on

How to remove a column from a DataFrame, with some extra detail

Removing one or more columns from a pandas DataFrame is a pretty common task, but it turns out there are a number of possible ways to perform this task. I found that this StackOverflow question, along with solutions and discussion in it raised a number of interesting topics. It is worth digging in a little bit to the details.

First, what’s the “correct” way to remove a column from a DataFrame? The standard way to do this is to think in SQL and use drop.

import pandas as pd
import numpy as np

df = pd.DataFrame(np.arange(25).reshape((5,5)),               
                  columns=list("abcde"))

display(df)

try:
    df.drop('b')
except KeyError as ke:
    print(ke)

   a  b  c  d  e
0  0  1  2  3  4
1  5  6  7  8  9
2 10 11 12 13 14
3 15 16 17 18 19
4 20 21 22 23 24
"['b'] not found in axis"
Enter fullscreen mode Exit fullscreen mode

Wait, what? Why an error? That’s because the default axis that drop works with is the rows. As with many pandas methods, there’s more than one way to invoke the method (which some people find frustrating).

You can drop rows using axis=0 or axis='rows', or using the labels argument.

df.drop(0) # drop a row, on axis 0 or 'rows'
df.drop(0, axis=0) # same
df.drop(0, axis='rows') # same
df.drop(labels=0) # same
df.drop(labels=[0]) # same
Enter fullscreen mode Exit fullscreen mode
   a  b  c  d  e
1  5  6  7  8  9
2 10 11 12 13 14
3 15 16 17 18 19
4 20 21 22 23 24
Enter fullscreen mode Exit fullscreen mode

Again, how do we drop a column?

We want to drop a column, so what does that look like? You can specify the axis or use the columns parameter.

df.drop('b', axis=1) # drop a column
df.drop('b', axis='columns') # same
df.drop(columns='b') # same
df.drop(columns=['b']) # same
Enter fullscreen mode Exit fullscreen mode
   a  c  d  e
0  0  2  3  4
1  5  7  8  9
2 10 12 13 14
3 15 17 18 19
4 20 22 23 24
Enter fullscreen mode Exit fullscreen mode

There you go, that’s how you drop a column. Now you have to either assign to a new variable, or back to your old variable, or pass in inplace=True to make the change permanent.

df2 = df.drop('b', axis=1)

print(df2.columns)
print(df.columns)

Index(['a', 'c', 'd', 'e'], dtype='object')
Index(['a', 'b', 'c', 'd', 'e'], dtype='object')
Enter fullscreen mode Exit fullscreen mode

It’s also worth noting that you can drop both rows and columns at the same time using drop by using the index and columns arguments at once, and you can pass in multiple values.

df.drop(index=[0,2], columns=['b','c'])
Enter fullscreen mode Exit fullscreen mode
   a  d  e
1  5  8  9
3 15 18 19
4 20 23 24
Enter fullscreen mode Exit fullscreen mode

If you didn’t have the drop method, you can basically obtain the same results through indexing. There are many ways to accomplish this, but one equivalent solution is indexing using the .loc indexer and isin, along with inverting the selection.

df.loc[~df.index.isin([0,2]), ~df.columns.isin(['b', 'c'])]
Enter fullscreen mode Exit fullscreen mode
   a  d  e
1  5  8  9
3 15 18 19
4 20 23 24
Enter fullscreen mode Exit fullscreen mode

If none of that makes sense to you, I would suggest reading through my series on selecting and indexing in pandas, starting here.

Back to the question

Looking back at the original question though, we see there is another available technique for removing a column.

del df['a']
df
Enter fullscreen mode Exit fullscreen mode
   b  c  d  e
0  1  2  3  4
1  6  7  8  9
2 11 12 13 14
3 16 17 18 19
4 21 22 23 24
Enter fullscreen mode Exit fullscreen mode

Poof! It’s gone. This is like doing a drop with inplace=True.

What about attribute access?

We also know that we can use attribute access to select columns of a DataFrame.

df.b
Enter fullscreen mode Exit fullscreen mode
0  1
1  6
2 11
3 16
4 21
Name: b, dtype: int64
Enter fullscreen mode Exit fullscreen mode

Can we delete the column this way?

del df.b
Enter fullscreen mode Exit fullscreen mode
--------------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
<ipython-input-10-0dca358a6ef9> in <module>
---------> 1 del df.b

AttributeError: b
Enter fullscreen mode Exit fullscreen mode

We cannot. This is not an option for removing columns with the current pandas design. Is this technically impossible? How come del df['b']works but del df.b doesn’t?. Let’s dig into those details and see whether it would be possible to make the second work as well.

The first version works because in pandas, the DataFrame implements the __delitem__ method which gets invoked when you execute del df['b']. But what about del df.b, is there a way to handle that?

First, let’s make a simple class that shows how this works under the hood. Instead of being a real DataFrame, we’ll just use a dict as a container for our columns (which could really contain anything, we’re not doing any indexing here).

class StupidFrame:
    def __init__ (self, columns):
        self.columns = columns

    def __delitem__ (self, item):
        del self.columns[item]

    def __getitem__ (self, item):
        return self.columns[item]

    def __setitem__ (self, item, val):
        self.columns[item] = val

f = StupidFrame({'a': 1, 'b': 2, 'c': 3})
print("StupidFrame value for a:", f['a'])
print("StupidFrame columns: ", f.columns)
del f['b']
f.d = 4
print("StupidFrame columns: ", f.columns)
Enter fullscreen mode Exit fullscreen mode
StupidFrame value for a: 1
StupidFrame columns: {'a': 1, 'b': 2, 'c': 3}
StupidFrame columns: {'a': 1, 'c': 3}
Enter fullscreen mode Exit fullscreen mode

A couple of things to note here. First, we how that we can access the data in our StupidFrame with the index operators ([]), and use that for setting, getting, and deleting items. When we assigned d to our frame, it wasn’t added to our columns because it’s just a normal instance attribute. If we wanted to be able to handle the columns as attributes, we have to do a little bit more work.

So following the example from pandas (which supports attribute access of columns), we add the __getattr__ method, but we also will handle setting it with the __setattr__ method and pretend that any attribute assignment is a ‘column’. We have to update our instance dictionary (__dict__) directly to avoid an infinite recursion.

class StupidFrameAttr:
    def __init__ (self, columns):
        self. __dict__ ['columns'] = columns

    def __delitem__ (self, item):
        del self. __dict__ ['columns'][item]

    def __getitem__ (self, item):
        return self. __dict__ ['columns'][item]

    def __setitem__ (self, item, val):
        self. __dict__ ['columns'][item] = val

    def __getattr__ (self, item):
        if item in self. __dict__ ['columns']:
            return self. __dict__ ['columns'][item]
        elif item == 'columns':
            return self. __dict__ [item]
        else:
            raise AttributeError

    def __setattr__ (self, item, val):
        if item != 'columns':
            self. __dict__ ['columns'][item] = val
        else:
            raise ValueError("Overwriting columns prohibited") 


f = StupidFrameAttr({'a': 1, 'b': 2, 'c': 3})
print("StupidFrameAttr value for a", f['a'])
print("StupidFrameAttr columns: ", f.columns)
del f['b']
print("StupidFrameAttr columns: ", f.columns)
print("StupidFrameAttr value for a", f.a)
f.d = 4
print("StupidFrameAttr columns: ", f.columns)
del f['d']
print("StupidFrameAttr columns: ", f.columns)
f.d = 5
print("StupidFrameAttr columns: ", f.columns)
del f.d
Enter fullscreen mode Exit fullscreen mode
StupidFrameAttr value for a 1
StupidFrameAttr columns: {'a': 1, 'b': 2, 'c': 3}
StupidFrameAttr columns: {'a': 1, 'c': 3}
StupidFrameAttr value for a 1
StupidFrameAttr columns: {'a': 1, 'c': 3, 'd': 4}
StupidFrameAttr columns: {'a': 1, 'c': 3}
StupidFrameAttr columns: {'a': 1, 'c': 3, 'd': 5}
--------------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
<ipython-input-12-fd29f59ea01e> in <module>
     39 f.d = 5
     40 print("StupidFrameAttr columns: ", f.columns)
--------> 41 del f.d

AttributeError: d
Enter fullscreen mode Exit fullscreen mode

How could we handle deletion?

Everything works but deletion using attribute access. We handle setting/getting columns using both the array index operator ([]) and attribute access. But what about detecting deletion? Is that possible?

One way to do this is using the __delattr__ method, which is described in the data model documentation. If you define this method in your class, it will be invoked instead of updating an instance’s attribute dictionary directly. This gives us a chance to redirect this to our columns instance.

class StupidFrameDelAttr(StupidFrameAttr):
    def __delattr__ (self, item):
        # trivial implementation using the data model methods
        del self. __dict__ ['columns'][item]

f = StupidFrameDelAttr({'a': 1, 'b': 2, 'c': 3})
print("StupidFrameDelAttr value for a", f['a'])
print("StupidFrameDelAttr columns: ", f.columns)
del f['b']
print("StupidFrameDelAttr columns: ", f.columns)
print("StupidFrameDelAttr value for a", f.a)
f.d = 4
print("StupidFrameDelAttr columns: ", f.columns)
del f.d 
print("StupidFrameDelAttr columns: ", f.columns)

Enter fullscreen mode Exit fullscreen mode
StupidFrameDelAttr value for a 1
StupidFrameDelAttr columns: {'a': 1, 'b': 2, 'c': 3}
StupidFrameDelAttr columns: {'a': 1, 'c': 3}
StupidFrameDelAttr value for a 1
StupidFrameDelAttr columns: {'a': 1, 'c': 3, 'd': 4}
StupidFrameDelAttr columns: {'a': 1, 'c': 3}
Enter fullscreen mode Exit fullscreen mode

Now I’m not suggesting that attribute deletion for columns would be easy to add to pandas, but at least this shows how it could be possible. In the case of current pandas, deleting columns is best done using drop.

Also, it’s worth mentioning here that when you create a new column in pandas, you don’t assign it as an attribute. To better understand how to properly create a column, you can check out this article.

If you already knew how to drop a column in pandas, hopefully you understand a little bit more about how this works.

The post How to remove a column from a DataFrame, with some extra detail appeared first on wrighters.io.

Discussion (0)