提高For Loop Python的速度

Question

I have a function that returns a dictionary. The function works by calculating values based on an array in a dataframe.

The dataframe has about 1000,000 rows and looks like this:

                  col1                  
row1         [2, 3, 44, 89.6,...]           
row2         [10, 4, 33.3, 1.11,...]
row3         [3, 4, 3, 2.6, 5.9, 8, 10,...]  

My function takes in each array in each row, does some calculations and returns a dictionary based on these calculations. However, it is very slow. There is a lot of data to sift through which I appreciate but is there a way that I can improve the speed?

The Issues Dataframe is long. Each array can contain 100+ values. Ranges from about 10-80.

My code looks like this:

list1 = []

for i in df.itertuples():
    list1.append(list(function(i.data).values()))

The idea here is that I loop through each row in 'df', apply my function to the 'data' column and append the results to a list, 'list1'.

Function Explained

My function computes some pretty basic stuff. It takes in an array as a parameter and calculates stuff based on that array eg how long is it, average value in the array, min and max of array. I compute 8 values and store them in a dictionary. The last thing my function does is look at these computed values and add a final key to the dictionary in the form of a boolean.

Answer 1

Like I said in the comments, if your function is costly (reducing each row is the time consuming part of your code), then a first step is to use multiprocessing because it's easy to test.

Here is something you could try:

import time
from multiprocessing import Pool

def f(x):
  time.sleep(10*10**-6) # Faking complex computation
  return x

def seq_test(input_array):
  return list(map(f, input_array))

def par_test(input_array):
  pool = Pool(8)  #  "nproc --all" or "sysctl -n hw.ncpu" on osx
  return pool.map(f, input_array)

def run_test(test_function):
  test_size = 10*10**4
  test_input = [i for i in range(test_size)]

  t0 = time.time()
  result = test_function(test_input)
  t1 = time.time()

  print(f"{test_function.__name__}: {t1-t0:.3f}s")

run_test(seq_test)
run_test(par_test)

On my machine the parallel version runs about 7 times faster (quite close to the factor 8 we could hope for):

seq_test: 2.131s
par_test: 0.300s

If that's not enough, the next step is to write function f in a different language, once again, what seems simpler here is to go for Cython . But for discussing that we need to see what's inside your function.

Answer 2

I suggest change format of your data like:

print (df)
                            col1
row1            [2, 3, 44, 89.6]
row2         [10, 4, 33.3, 1.11]
row3  [3, 4, 3, 2.6, 5.9, 8, 10]

from itertools import chain

df = pd.DataFrame({
    'idx' : df.index.repeat(df['col1'].str.len()),
    'col1' : list(chain.from_iterable(df['col1'].tolist()))
})
print (df)
     idx   col1
0   row1   2.00
1   row1   3.00
2   row1  44.00
3   row1  89.60
4   row2  10.00
5   row2   4.00
6   row2  33.30
7   row2   1.11
8   row3   3.00
9   row3   4.00
10  row3   3.00
11  row3   2.60
12  row3   5.90
13  row3   8.00
14  row3  10.00

And then aggregate your data:

df1 = df.groupby('idx')['col1'].agg(['sum','mean','max','min'])
print (df1)
         sum       mean   max   min
idx                                
row1  138.60  34.650000  89.6  2.00
row2   48.41  12.102500  33.3  1.11
row3   36.50   5.214286  10.0  2.60