How to iterate over strings in pandas dataframe and remove unwanted words?

Viewed 321

I'm trying to clean up a Pandas dataframe called names2. It consists of 599,864 rows, 549,317 of them are non-null. In each row under the column in question, 'primary_profession' there is either 1 string, an array of strings or NaN.

Here is a look at how I loaded the dataframe:

name_basics_imdb = pd.read_csv('imdb.name.basics.csv.gz')
names = name_basics_imdb
names2 = names.copy(deep=True)

(Note: I dropped some columns and rows and renamed a column, if you need more details they'll gladly be supplied)

Here is a view of names2.info( )

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 599865 entries, 0 to 599864
Data columns (total 4 columns):
 #   Column              Non-Null Count   Dtype 
---  ------              --------------   ----- 
 0   nconst              599865 non-null  object
 1   primary_name        599865 non-null  object
 2   primary_profession  549317 non-null  object
 3   known_for_titles    569766 non-null  object
dtypes: object(4)
memory usage: 18.3+ MB

names2.head( )

    nconst       primary_name                  primary_profession                         known_for_titles
0   nm0061671   Mary Ellen Bauder   miscellaneous,production_manager,producer   tt0837562,tt2398241,tt0844471,tt0118553
1   nm0061865   Joseph Bauer       composer,music_department,sound_department   tt0896534,tt6791238,tt0287072,tt1682940
2   nm0062070   Bruce Baum                miscellaneous,actor,writer            tt0363631
3   nm0062195   Axel Baumann    camera_department,cinematographer,art_department    tt0114371,tt2004304,tt1618448,tt1224387
4   nm0062798   Pete Baxter production_designer,art_department,set_decorator    tt0452644,tt0452692,tt34580

The goal is to iterate through each row's string, strings, or Nan and just keep the rows with writer, writer director, or just director. Any other profession can be thrown out. For instance, in row 2: miscellaneous, actor, writer are in the primary_profession column, miscellaneous and actor can be eliminated leaving only writer in that row.

Any rows where there is no writer or director or contain NaN can be dropped.

Here are a few of the attempts I made

#inverse filtering

value_list = ['miscellaneous', 'production_manager', 'composer', 'music_department', 'sound_department',
      'miscellaneous', 'actor', '...', 'costume_department', 'costume_designer', 'actress', 'art_director', 'music_department' ]

#have to split the arrays first

inverse_bool_series = ~names2.primary_profession.isin(value_list)
names2_filtered = names2[inverse_bool_series]
names2_filtered

I also tried

names2['primary_profession'] = names2['primary_profession'].str.split(",").str[:3]
names2['primary_profession']
(names2['primary_profession'][0])
type(names2['primary_profession'][0][0]) 

And then there was this

for index, row in names2.iterrows():
    idx = list(len(range(names2.primary_profession)))
    for i in idx:
        print(row['primary_profession'][i])

To summarize, the goal is for the dataframe names2 to have only rows with the profession writer, writer director, or just writer

3 Answers
In [96]: names2
Out[96]:
      nconst       primary_name                                         primary_profession                         known_for_titles
0  nm0061671  Mary Ellen Bauder  miscellaneous,production_manager,producer,director,writer  tt0837562,tt2398241,tt0844471,tt0118553
1  nm0061865       Joseph Bauer                 composer,music_department,sound_department  tt0896534,tt6791238,tt0287072,tt1682940
2  nm0062070         Bruce Baum                                 miscellaneous,actor,writer                                tt0363631
3  nm0062195       Axel Baumann           camera_department,cinematographer,art_department  tt0114371,tt2004304,tt1618448,tt1224387
4  nm0062798        Pete Baxter           production_designer,art_department,set_decorator              tt0452644,tt0452692,tt34580

In [97]: profs = names2['primary_profession'].str.split(',').explode()

In [98]: profs
Out[98]:
0          miscellaneous
0     production_manager
0               producer
0               director
0                 writer
1               composer
1       music_department
1       sound_department
2          miscellaneous
2                  actor
2                 writer
3      camera_department
3        cinematographer
3         art_department
4    production_designer
4         art_department
4          set_decorator
Name: primary_profession, dtype: object

In [99]: filtered_profs = profs[profs.isin(['writer', 'writer director', 'director'])]

In [100]: filtered_profs.groupby(filtered_profs.index).agg(','.join)
Out[100]:
0    director,writer
2             writer
Name: primary_profession, dtype: object

In [101]: names2.drop('primary_profession', axis=1).join(filtered_profs.groupby(filtered_profs.index).agg(','.join), how='inner')
Out[101]:
      nconst       primary_name                         known_for_titles primary_profession
0  nm0061671  Mary Ellen Bauder  tt0837562,tt2398241,tt0844471,tt0118553    director,writer
2  nm0062070         Bruce Baum                                tt0363631             writer
# set the words you want to match.
matched_words = ['writer', 'writer_director', 'director']

#drop rows which has nan in column 'primary_profession'
names2 = names.dropna(axis='index', subset=['primary_profession'])

#extract all matched words
names2_extractall = names2['primary_profession'].str.extractall(rf'({"|".join(matched_words)})')

#groupby index and join those matches result by ','
mod_prof = names2_extractall.groupby(level=0).apply(lambda x: ",".join(x.iloc[:, 0]))

#assign to column 'primary_profession'
names2 = names2.assign(primary_profession=mod_prof)

#drop no matched rows
names2 = names2.dropna(axis='index', subset=['primary_profession'])
unwanted_row_indices=[]
for index,row in names2.iterrows():
    if('writer' in row['primary_profession'].lower()):
        pass
    else:
        unwanted_row_indices.append(index)
names2=names2.drop(unwanted_row_indices,axis=0)
Related