Speed up simple Python function that uses list comprehension

Question

I'm extracting 4 columns from an imported CSV file (~500MB) to be used for fitting a scikit-learn regression model.

It seems that this function used to do the extraction is extremely slow. I just learnt python today, any suggestions on how the function can be sped up?

Can multithreading/core be used? My system has 4 cores.

def splitData(jobs):
    salaries = [jobs[i]['salaryNormalized'] for i, v in enumerate(jobs)]
    descriptions = [jobs[i]['description'] + jobs[i]['normalizedLocation'] + jobs[i]['category'] for i, v in enumerate(jobs)]
    titles = [jobs[i]['title'] for i, v in enumerate(jobs)]

    return salaries, descriptions, titles

---

Full Code

def loadData(filePath):
    reader = csv.reader( open(filePath) )
    rows = []

    for i, row in enumerate(reader):
        categories = ["id", "title", "description", "rawLocation", "normalizedLocation",
                        "contractType", "contractTime", "company", "category",
                        "salaryRaw", "salaryNormalized","sourceName"]

        # Skip header row
        if i != 0: 
            rows.append( dict(zip(categories, row)) )

    return rows



def splitData(jobs):
    salaries = []
    descriptions = []
    titles = []

    for i in xrange(len(jobs)):
        salaries.append( jobs[i]['salaryNormalized'] )
        descriptions.append( jobs[i]['description'] + jobs[i]['normalizedLocation'] + jobs[i]['category'] )
        titles.append( jobs[i]['title'] )

    return salaries, descriptions, titles



def fit(salaries, descriptions, titles):
    #Vectorize
    vect = TfidfVectorizer()
    vect2 = TfidfVectorizer()
    descriptions = vect.fit_transform(descriptions)
    titles = vect2.fit_transform(titles)

    #Fit
    X = hstack((descriptions, titles))
    y = [ np.log(float(salaries[i])) for i, v in enumerate(salaries) ]

    rr = Ridge(alpha=0.035)
    rr.fit(X, y)

    return vect, vect2, rr, X, y



jobs = loadData( paths['train_data_path'] )
salaries, descriptions, titles = splitData(jobs)
vect, vect2, rr, X_train, y_train = fit(salaries, descriptions, titles)

Answer 1

def loadData(filePath):
    reader = csv.reader( open(filePath) )
    rows = []

    for i, row in enumerate(reader):
        categories = ["id", "title", "description", "rawLocation", "normalizedLocation",
                        "contractType", "contractTime", "company", "category",
                        "salaryRaw", "salaryNormalized","sourceName"]

A list like this should really be a global variable to avoid the expense of recreating it constantly. But you'd do better not store your stuff in a dictionary. Instead do this:

   for (id, title, description, raw_location, normalized_location, contract_type,
       contractTime, company, category, salary_raw, salary_normalized, source_name) in reader:

       yield salary_normalized, ''.join((description, normalized_location, category)), title

All the stuff is then stored in python local variables (fairly efficient). Then yield produces the three elements you are actually wanting. Just use

   salaries, descriptions, titles = zip(*loadData(...))

to get your three lists again

        # Skip header row
        if i != 0: 
            rows.append( dict(zip(categories, row)) )

Rather than this, call reader.next() before the loop to take out the header

    return rows

Answer 2

titles = [jobs[i]['title'] for i, v in enumerate(jobs)] can (should?) be rewritten :

titles = [j['title'] for j in jobs.items()] because we just want to access the value at position i (More details)

Thus, the whole code would be :

def splitData(jobs):
    salaries = [j['salaryNormalized'] for j in jobs.items)]
    descriptions = [j['description'] + j['normalizedLocation'] + j['category'] for j in jobs.items)]
    titles = [j['title'] for j in jobs.items)]

    return salaries, descriptions, titles

Not quite sure how much it helps from a performance point of view.

Edit : Otherwise, another option might be to write a generator which returns j['salaryNormalized'], j['description'] + j['normalizedLocation'] + j['category'], j['title'] as you need it. It depends how you use your function really.

Answer 3

It seems that you are looping through the entire job data set 3 times (once each for salaries, descriptions and titles). You will be able to speed this up three-fold, if you extract all the info in one pass:

def split_data(jobs):
    for job, info in jobs.items():
        salaries.append(info['salaryNormalized'])
        descriptions.append([info['description'], 
                             info['normalizedLocation'], 
                             info['category']])
        titles.append(info['title'])

EDIT Added loadData(); slight tweak to return a dictionary of dictionaries, instead of a list of dictionaries:

def load_data(filepath):
    reader = csv.reader(open(filepath))
    jobs = {}

    for i, row in enumerate(reader):
        categories = ["id", "title", "description", "rawLocation", "normalizedLocation",
                        "contractType", "contractTime", "company", "category",
                        "salaryRaw", "salaryNormalized","sourceName"]

        if i != 0:
            jobs[i] = dict(zip(categories, row))

    return jobs

An example:

jobs = {0 : {'salaryNormalized' : 10000, 'description' : 'myKob',
             'normalizedLocation': 'Hawaii', 'category': 'categ1',
             'title' : 'tourist'},
        1 : {'salaryNormalized' : 15000, 'description' : 'myKob',
             'normalizedLocation': 'Hawaii', 'category': 'categ2',
             'title' : 'tourist'},
        2 : {'salaryNormalized' : 50000, 'description' : 'myKob',
             'normalizedLocation': 'Hawaii', 'category': 'categ10',
             'title' : 'resort_manager'}}

salaries, descriptions, titles = [], [], []
split_data(jobs)

print(salaries)
--> [10000, 15000, 50000]

print(descriptions)
--> [['myKob', 'Hawaii', 'categ1'], ['myKob', 'Hawaii', 'categ2'], ['myKob', 'Hawaii', 'categ10']]

print(titles)
--> ['tourist', 'tourist', 'resort_manager']

Hope this helps!