加快大型 arrays 和数据集的操作（熊猫慢，Numpy 更好，进一步改进？）

Question

I have a large dataset comprising millions of rows and around 6 columns. The data is currently in a Pandas dataframe and I'm looking for the fastest way to operate on it. For example, let's say I want to drop all the rows where the value in one column is "1".

Here's my minimal working example:

# Create dummy data arrays and pandas dataframe
array_size = int(5e6)
array1 = np.random.rand(array_size)
array2 = np.random.rand(array_size)
array3 = np.random.rand(array_size)
array_condition = np.random.randint(0, 3, size=array_size)

df = pd.DataFrame({'array_condition': array_condition, 'array1': array1, 'array2': array2, 'array3': array3})

def method1():
    df_new = df.drop(df[df.array_condition == 1].index)

EDIT: As Henry Yik pointed out in the comments, a faster Pandas approach is this:

def method1b():
    df_new = df[df.array_condition != 1]

I believe that Pandas can be quite slow at this sort of thing, so I also implemented a method using numpy, processing each column as a separate array:

def method2():
    masking = array_condition != 1
    array1_new = array1[masking]
    array2_new = array2[masking]
    array3_new = array3[masking]
    array_condition_new = array_condition[masking]    

And the results:

%timeit method1()
625 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit methodb()
158 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit method2()
138 ms ± 3.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

So we do see a slight significant performance boost using numpy. However, this is at the cost of much less readable code (ie having to create a mask and apply it to each array). This method doesn't seem as scalable either as if I have, say, 30 columns of data, I'll need a lot of lines of code that apply the mask to every array, Additionally, it would be useful to allow optional columns. so this method may fail trying to operate on arrays which are empty.

Therefore, I have 2 questions:

1) Is there a cleaner / more flexible way to implement this in numpy?

2) Or better, is there any higher performance method I could use here? eg JIT (numba?), Cython or something else?

PS, in practice, in-place operations can be used, replacing the old array with the new one once data is dropped

Answer 1

You may find using numpy.where useful here. It converts a Boolean mask to array indices, making life much cheaper. Combining this with numpy.vstack allows for some memory-cheap operations:

def method3():
    wh = np.where(array_condition == 1)
    return np.vstack(tuple(col[wh] for col in (array1, array2, array3)))

This gives the following timeits:

>>> %timeit method2()
180 ms ± 6.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit method3()
96.9 ms ± 2.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Tuple unpacking allows the operation to be fairly light on memory, as when the object is vstack-ed back together, it is smaller. If you need to get your columns out of a DataFrame directly, the following code snippet may be useful:

def method3b():
    wh = np.where(array_condition == 1)
    col_names = ['array1','array2','array3']
    return np.vstack(tuple(col[wh] for col in tuple(df[col_name].to_numpy() 
        for col_name in col_names)))

This allows one to grab columns by name from the DataFrame, which are then tuple unpacked on the fly. The speed is about the same:

>>> %timeit method3b()
96.6 ms ± 3.09 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Enjoy!

Answer 2

Part 1: Pandas and (maybe) Numpy

Compare your method1b and method2 :

• method1b generates a DataFrame , which is probably what you want,
• method2 generates a Numpy array , so to get fully comparable result, you should subsequently generate a DataFrame from it.

So I changed your method2 to:

def method2():
    masking = array_condition != 1
    array1_new = array1[masking]
    array2_new = array2[masking]
    array3_new = array3[masking]
    array_condition_new = array_condition[masking]
    df_new = pd.DataFrame({ 'array_condition': array_condition[masking],
        'array1': array1_new, 'array2': array2_new, 'array3': array3_new})

and then compared execution times (using %timeit ).

The result was that my (expanded) version of method2 executed about 5% longer than method1b (check on your own).

So my opinion is that as long as a single operation is concerned, it is probably better to stay with Pandas .

But if you want to perform on your source DataFrame a couple of operations in sequence and / or you are satisfied with the result as a Numpy array, it is worth to:

• Call arr = df.values to get the underlying Numpy array.
• Perform all required operations on it using Numpy methods.
• (Optionally) create a DataFrame from the final reslut.

I tried Numpy version of method1b :

def method3():
    a = df.values
    arr = a[a[:,0] != 1]

but the execution time was about 40 % longer .

The reason is probably that Numpy array has all elements of the same type, so array_condition column is coerced to float and then the whole Numpy array is created, what takes some time.

Part 2: Numpy and Numba

An alternative to consider is to use Numba package - a Just-In-Time Python compiler.

I made such test:

Created a Numpy array (as a preliminary step):

a = df.values

The reason is that JIT compiled methods are able to use Numpy methods and types, but not those of Pandas .

To perform the test, I used almost the same method as above, but with @njit annotation (requires from numba import njit ):

@njit
def method4():
    arr = a[a[:,0] != 1]

This time:

• The execution time was about 45 % of the time for method1b .
• But since a = df.values has been executed before the test loop, there are doubts whether this result is comparable with earlier tests.

Anyway, try Numba on your own, maybe it will be an interesting option for you.