An Easy Way to Divide Your Dataset Based on Data Types with Pandas | by Byron Dolon | Towards Data Science

Before you can get into pre-processing data for a machine learning model to use, you’ll usually need to separate your numerical and non-numerical columns. In this piece, we’ll go over an easy way to do just that!
One way to do it is to manually check which columns contain just numbers and which columns contain text. However, if you’re working with a dataset that has dozens of columns, having to write code like the line below is a little silly and prone to human error.
df_numerical = df[['num_1', 'num_2', 'num_3', 'num_4', 'num_5']]
Instead, let’s take a look at using a simple Pandas method that will do this without you having to manually type out all your desired columns.
If you’d like to follow along, go ahead and download the ‘train’ dataset here. Once you’ve done that, make sure it’s saved to the same directory as your notebook and then run the code below to read it in:
import pandas as pddf = pd.read_csv('train.csv')
Now let’s get going!
A good first step before we split our DataFrame’s columns is to check what columns we’re actually working with, as well as what their respective data types are. To do so, we use the following Pandas method:
df.info()
Using info gives us a breakdown of all of the columns in our DataFrame, how many non-null values each column contains, and the DataTypes. We verify that there are indeed text columns in this dataset, which in Pandas fall under the “object” data type. As such, before we can actually train a machine learning model on this data set, we will need to somehow convert these text columns into numerical columns. A simple way to do that would be to pick an encoding method and apply it to all categorical columns simultaneously.
Therefore, let’s separate our numerical and categorical columns using the select_dtypes method in Pandas. We’ll run the following two lines of code and get the resulting subsets of columns:
df_numerical = df.select_dtypes(include='int64')
df_categorical = df.select_dtypes(include='object')
Wonderful!
The select_dtypes method returns a subset of the DataFrame based on the column data types you specify. In this case, we’ve used the include parameter and passed in 'int64' to get the numerical columns and 'object' to get the categorical columns.
BUT... I just realized I actually did something wrong.
The subset of categorical columns looks correct, as it contains all the columns with text values in our original DataFrame. However, the subset of numerical columns seems to be missing a few columns. This is because we used select_dtypes to only get the 'int64' type columns, when in fact the info method earlier showed us that there are also 'float64' type columns. One way to get around this would be to pass a list into the 'include' parameter, but an alternative to that would be to use the 'exclude' parameter instead:
df_numerical = df.select_dtypes(exclude=’object’)
Now we have all 7 numerical columns from our original dataset! The exclude parameter returns all the columns except the data type specified, which in our case are the text (object) columns.
To actually get the column names and not only the subset DataFrames, we just need to do the following:
df_numerical_cols = df_numerical.columns.tolist()df_categorical_cols = df_categorical.columns.tolist()
Printing the two resulting lists gives us these two lists of column names:
Now that we’ve gotten the numerical and categorical columns separated, we can easily perform different transformations on our data before training a machine learning model. For example, if you wanted to impute (fill in missing data) your numerical and categorical columns with two different methods, you would be able to pass each list of column names to the respective imputation method (depending on what library you’re using).
I hope you found this quick run-through of select_dtypes helpful for splitting up your data and getting lists of columns you can pass into your pre-processing functions!
Good luck with your data-wrangling adventures!
More by me:- Conditional Selection and Assignment With .loc in Pandas- 2 Easy Ways to Get Tables From a Website With Pandas- 5 (and a half) Lines of Code for Understanding Your Data with Pandas- Top 4 Repositories on GitHub to Learn Pandas- Learning to Forecast With Tableau in 5 Minutes Or Less
*sidenote- I know the target (column with the value you need to predict if you were to train a machine learning model) is still in the df_numerical list (woops), so if you do use this code be sure to assign the “Survived” column to another variable and then drop it from the DataFrame or list of column names.