3
\$\begingroup\$

Aim of the Code:

The aim of the function I'm working on is to take a particular column from a DataFrame that contains multiple <tag><content> pairs in the form of a string and expand them into individual columns in an efficient way.

An example file can be found here (aprox. 3M entries), which can be loaded as a DataFrame with:

def load_refseq_coordinates(ifile):
    # https://m.ensembl.org/info/website/upload/gff3.html
    return pd.read_csv(ifile, sep='\t', comment='#',
                       names=['seqid', 'source', 'type', 'start', 'end', 'score',
                              'strand', 'phase', 'attributes'])
    return df

The loaded DataFrame will have a column named attributes that has content such as:

attributes
ID=NC_000001.11:1..248956422;Dbxref=taxon:9606;Name=1;chromosome=1;gbkey=Src;genome=chromosome;mol_type=genomic DNA

with different pairs separated by ; and <tag><content> pairs defined with =.

The aim would, then to keep the rest of the columns of the DataFrame and get this one converted to:

attributes.ID attributes.Dbxref attributes.Name attributes.chromosome attributes.gbkey attributes.genome attributes.mol_type
NC_000001.11:1..248956422 taxon:9606 1 1 Src chromosome genomic DNA

And the same for all rows of the DataFrame.

The Issue:

All the solutions I've found so far do not scale proficiently and applying them to the 3M entry file takes just too long.

In the last 3 cases I'll show, performance seems better, as I can test up to 100K in around 4s (with that time is around 10K rows in the first solutions), but it is still poor performance for the 3M rows.

Current Solutions:

Note that:

  • I'm running this in a Jupyter Notebook, thus the %timeit call.
  • I'm only checking up to 10k rows to see how it scales (gets too long with more), except the last 3 solutions, as they perform better (thus I time them for 100K rows)

1. Series.apply a function that does the double string split

def expand_info_string1(df, column='info', entity_sep=';', id_sep='='):
    def do_expand(cell, entity_seq, id_sep):
        data = [x.split(id_sep) for x in str(cell).split(entity_seq)]
        return pd.Series(list(zip(*data))[-1], index=list(zip(*data))[0])
        
    info = df[column].apply(do_expand, entity_seq=entity_sep, id_sep=id_sep)
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1).fillna('')

%timeit dd = expand_info_string1(df.sample(10), 'attributes')
%timeit dd = expand_info_string1(df.sample(100), 'attributes')
%timeit dd = expand_info_string1(df.sample(1000), 'attributes')
%timeit dd = expand_info_string1(df.sample(10000), 'attributes')

330 ms ± 267 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
299 ms ± 18.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
708 ms ± 375 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.4 s ± 263 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

2. Series.str.split + Series.apply a function to do 1 split

def expand_info_string3(df, column='info', entity_sep=';', id_sep='='):
    def do_expand(cell, id_sep):
        data = [str(x).split(id_sep) for x in cell]
        return pd.Series(list(zip(*data))[-1], index=list(zip(*data))[0])
        
    info = df[column].str.split(entity_sep).apply(do_expand, id_sep=id_sep)
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1).fillna('')

%timeit dd = expand_info_string3(df.sample(10), 'attributes')
%timeit dd = expand_info_string3(df.sample(100), 'attributes')
%timeit dd = expand_info_string3(df.sample(1000), 'attributes')
%timeit dd = expand_info_string3(df.sample(10000), 'attributes')

90 ms ± 28.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
338 ms ± 66.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
553 ms ± 13.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.55 s ± 379 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

3. Series.str.split(expand=True) + Series.apply a function to do 1 split

def expand_info_string2(df, column='info', entity_sep=';', id_sep='='):
    def do_expand(cell):
        data = cell.dropna().values
        return pd.Series(list(zip(*data))[-1], index=list(zip(*data))[0])
        
    info = df[column].str.split(entity_sep, expand=True).apply(lambda cell: cell.str.split(id_sep), axis=1).apply(do_expand, axis=1)
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1).fillna('')

%timeit dd = expand_info_string2(df.sample(10), 'attributes')
%timeit dd = expand_info_string2(df.sample(100), 'attributes')
%timeit dd = expand_info_string2(df.sample(1000), 'attributes')
%timeit dd = expand_info_string2(df.sample(10000), 'attributes')

236 ms ± 44.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
262 ms ± 19.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.35 s ± 281 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
11.2 s ± 313 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

4. Series.str.extractall + regex

def expand_info_regex(df, column='info', entity_sep=';', id_sep='='):
    
    info_re = re.compile("([^;]+?)=(?:([^;]+))?") 
    info= (df[column].str.extractall(info_re).reset_index(level=1, drop=True)
                     .set_index(0, append=True)[1]
                     .unstack(level=1))
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1).fillna('')
%timeit dd = expand_info_regex(df.sample(10), 'attributes')
%timeit dd = expand_info_regex(df.sample(100), 'attributes')
%timeit dd = expand_info_regex(df.sample(1000), 'attributes')
%timeit dd = expand_info_regex(df.sample(10000), 'attributes')
%timeit dd = expand_info_regex(df.sample(100000), 'attributes')

122 ms ± 3.41 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
129 ms ± 3.38 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
157 ms ± 3.03 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
474 ms ± 9.87 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.28 s ± 33.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

5. Series.str.split(expand).stack + DataFrame.groupby

def expand_info_stackgroup(df, column='info', entity_sep=';', id_sep='='):
        
    info = (df[column].str.split(entity_sep, expand=True).stack()
                      .str.split(id_sep, expand=True).reset_index()
                      .drop(columns='level_1').groupby(['level_0', 0]).first()
                      .unstack(fill_value='').reset_index(drop=True))
    info.columns = [x[1] for x in info.columns]
    info.index = df.index
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1)

%timeit dd = expand_info_stackgroup(df.sample(10), 'attributes')
%timeit dd = expand_info_stackgroup(df.sample(100), 'attributes')
%timeit dd = expand_info_stackgroup(df.sample(1000), 'attributes')
%timeit dd = expand_info_stackgroup(df.sample(10000), 'attributes')
%timeit dd = expand_info_stackgroup(df.sample(100000), 'attributes')

178 ms ± 24.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
269 ms ± 114 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
231 ms ± 87.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
494 ms ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.51 s ± 99.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

6. Series.str.split(expand).stack + DataFrame.set_index

def expand_info_stackindex(df, column='info', entity_sep=';', id_sep='='):

    info = (df[column].str.split(entity_sep, expand=True).stack()
                      .str.split(id_sep, expand=True).reset_index()
                      .drop(columns='level_1').set_index(['level_0', 0])
                  .unstack(fill_value='').reset_index(drop=True))
    info.columns = [x[1] for x in info.columns]
    info.index = df.index
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.')], axis=1)

%timeit dd = expand_info_stackindex(df.sample(10), 'attributes')
%timeit dd = expand_info_stackindex(df.sample(100), 'attributes')
%timeit dd = expand_info_stackindex(df.sample(1000), 'attributes')
%timeit dd = expand_info_stackindex(df.sample(10000), 'attributes')
%timeit dd = expand_info_stackindex(df.sample(100000), 'attributes')

200 ms ± 91.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
133 ms ± 4.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
158 ms ± 4.02 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
467 ms ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.16 s ± 252 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Any help in trying to optimise this so that is doable for file with high number of entries would be really appreciated. Thanks!

\$\endgroup\$
6
  • 1
    \$\begingroup\$ Incorporating advice from an answer into the question violates the question-and-answer nature of this site. You could post improved code as a new question, as an answer, or as a link to an external site - as described in I improved my code based on the reviews. What next?. I have rolled back the edit, so the answers make sense again. \$\endgroup\$
    – Toby Speight
    Commented Aug 6, 2021 at 14:55
  • \$\begingroup\$ Apologies. Fair enough. I'll add the modified function as a new answer, then. Thanks! \$\endgroup\$
    – jaumebonet
    Commented Aug 6, 2021 at 15:20
  • \$\begingroup\$ Or, if you want it to be reviewed and possibly further improved, post it as a new question. \$\endgroup\$
    – Toby Speight
    Commented Aug 6, 2021 at 15:22
  • 2
    \$\begingroup\$ I think I would just rather not repeat the question again, it would just look like a duplication. Posting it as a non-accepted self-answer seems to me more appropriate in this context. \$\endgroup\$
    – jaumebonet
    Commented Aug 6, 2021 at 15:29
  • \$\begingroup\$ Is it possible to obtain the source data in some other format which is faster to parse? Alternatively, have you considered something like fromfile from numpy? \$\endgroup\$
    – Juho
    Commented Sep 17, 2021 at 17:07

2 Answers 2

Reset to default
1
\$\begingroup\$

I've prepared a code snippet to solve your problem:

import pandas as pd
from tqdm import tqdm
import datetime as dt

def convert_data(input_file_path):
    print(f"converting file: {input_file_path}")
    output = []
    with open(input_file_path) as f:
        for i in tqdm(f):
            if i.startswith("#"):
                continue
            data = i.split("\t")[-1]
            item = dict()
            for j in data.split(";"):
                k, v = j.split("=")
                item[k] = v.strip()
            output.append(item)
    return output

def process_data(input_file_path):
    data = convert_data(input_file_path)
    print("converting to dataframe")
    return pd.DataFrame.from_dict(data)


start = dt.datetime.utcnow()
df = process_data("stackOverflow/pandas_split_columns_optimisation/test_data/data")
print(dt.datetime.utcnow() - start)
# converting file: stackOverflow/pandas_split_columns_optimisation/test_data/data
# 3854329it [00:28, 134941.12it/s]
# converting to dataframe
# 0:01:53.006208

Please check it, maybe it works for you.

\$\endgroup\$
1
  • \$\begingroup\$ Hi @sergei-malanin, thanks for your solution. It does not cover several points I was planning to solve (1) it does not keep the other columns, (2) it cannot be guided by column name, (3) it can only be applied on read, (4) it does not append the original column name to the new columns. That said, the solution is really fast in comparison to mine (0:02:57.497253 for the full 3.8M rows), I suspect because of the DataFrame.from_dict construction. Thus, I implemented that part into its own function following the others I have. Will update my question taking this adaptation into consideration. \$\endgroup\$
    – jaumebonet
    Commented Aug 6, 2021 at 13:40
0
\$\begingroup\$

@sergei-malanin's answer is really fast in comparison with the options I provided but does not completely meet the needs of the question, namely:

  1. it does not keep the other columns,
  2. it cannot be guided by column name,
  3. it can only be applied on read,
  4. it does not append the original column name to the new columns

As such, I took the idea of the DataFrame.from_dict construction of his solution and I adapted it to work within my requirements to showcase how it would look like:

def expand_info_fromdict(df, column='info', entity_sep=';', id_sep='='):
    def split_to_dict(data, entity_sep, id_sep):
        return dict([j.split(id_sep) for j in data.split(entity_sep)])

    info = pd.DataFrame.from_dict(list(df[column].apply(split_to_dict, entity_sep=entity_sep, id_sep=id_sep).values))
    info.index = df.index
    return pd.concat([df.drop(columns=[column]), info.add_prefix(f'{column}.').fillna('')], axis=1)

%timeit dd = expand_info_fromdict(df.sample(10), 'attributes')
%timeit dd = expand_info_fromdict(df.sample(100), 'attributes')
%timeit dd = expand_info_fromdict(df.sample(1000), 'attributes')
%timeit dd = expand_info_fromdict(df.sample(10000), 'attributes')
%timeit dd = expand_info_fromdict(df.sample(100000), 'attributes')

The times I got where somewhat better.

129 ms ± 30.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
175 ms ± 124 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
205 ms ± 23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
374 ms ± 6.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
2.42 s ± 45.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

For the 100k sample, this is almost a 40% improvement.

The full 3.8M rows run in just over 3 minutes.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.