Speed up simple Python function that uses list comprehension 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)
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