Numpy优化pandas数据框基于范围的合并的最快方式

在本文中，我们将介绍如何使用Numpy来优化pandas数据框基于范围的合并的速度。合并多个数据源是数据分析中常见的任务之一。对于对每个数据源的报告和分析，合并多个数据源是必要的，而在任何数据集中找到相关的数据点通常涉及到基于范围的或区间的合并。这就需要使用pandas数据框的merge操作。

阅读更多：Numpy 教程

问题描述

合并数据框时，基于范围的合并通常需要耗费很长时间，并且可能会导致崩溃。在大数据集上运行它时会带来挑战，因为每个数据源都需要被合并，而这可能会导致若干内存问题。输入数据通常是按一定规律进行分割的，因此基于范围的合并是优化方案。

解决方案

Numpy是一个高性能的Python库，专门用于数学计算。 Numpy数组本质上是大块的连续内存，它们的元素都是相同类型的。 Numpy数组可以比Python原生数据类型数组更快地处理大量数据，并提供了一个强大的库来处理大量计算问题。所以使用Numpy可以更快地完成任务，结合pandas数据框的merge方法可以更快地实现数据的基于范围的合并。

下面我们通过一个示例来具体说明。

假设我们要将一个数据集按照一组日期范围分割成多个数据源。我们将创建一个random_dates函数来生成一些随机日期，将其随机排列，然后将其转换为pandas的数据框。询问特定日期范围内的记录也是常见的需求。最快的方法是将日期范围划分为大小相等的块，然后手动合并这些块来提供结果。

import pandas as pd
import numpy as np

#创建一个随机的日期时间函数
def random_dates(start_date, end_date, n, unit='D', seed=None):
    if seed is not None:
        np.random.seed(seed)
    ndays = (end_date - start_date).days + 1
    return pd.to_timedelta(np.random.rand(n) * ndays, unit=unit) + start_date

#生成日期时间数据
start_date = pd.to_datetime('20210101')
end_date = pd.to_datetime('20211231')
n = 2000
df = pd.DataFrame({'Date': random_dates(start_date, end_date, n)})
df = df.sort_values(by='Date')

如上所述，我们创建了一个名为random_dates的函数来生成一些随机日期。然后，我们将其转换为pandas的数据框，并按日期排序。接下来，我们将数据集分成四个相等大小的块，以便进行从一个范围到另一个范围的查询。

#划分数据集
chunks = np.linspace(df.index.min(), df.index.max(), num=4, dtype=np.int)
for i in range(len(chunks) - 1):
    sub = df.loc[chunks[i]:chunks[i+1]]
    sub.to_csv(f'data/sub_{i}.csv', index=False)

此时，我们已将数据集划分为四个相等大小的数据集。

我们可以使用glob.glob检查我们在file_paths中保存的文件。

import glob
import os
file_paths = glob.glob('data/*.csv')
print(file_paths)

输出：

['data/sub_0.csv', 'data/sub_1.csv', 'data/sub_2.csv', 'data/sub_3.csv']

如果要在其日期范围内查询数据，最快的方法是将数据范围划分为大小相等的块，然后手动将这些块合并。

#查询数据
def query_data(start_date, end_date):
    dfs = []
    for file_path in file_paths:
        sub = pd.read_csv(file_path)
        sub_range = (sub['Date'] >= start_date) & (sub['Date'] <= end_date)
        sub = sub[sub_range]
        dfs.append(sub)
    result = pd.concat(dfs)
    return result.sort_values(by='Date')

如上所述，我们编写了一个query_data函数来查询指定日期范围内的数据。它遍历所有块，将符合条件的子数据框存储到列表中，最后使用pandas的concat方法合并这些数据框，并按日期排序。

优化方法

虽然我们已经编写了一个良好的查询函数来合并数据，但是对于大量数据，这个方法可能会变得很慢。这里我们介绍一个使用Numpy数组来优化这个过程的方法。

在查询日期范围时，我们可以使用Numpy数组来获得比使用pandas数据框更快的性能。我们可以使用Numpy的searchsorted函数在每个块中查找开始和结束时间的索引，以便确定在哪些块中需要查询数据。为此，我们需要将pandas的Series对象转换为Numpy数组。

def query_data_np(start_date, end_date):
    start_indices = []
    end_indices = []
    for file_path in file_paths:
        sub = pd.read_csv(file_path)['Date'].values
        start_idx = np.searchsorted(sub, start_date, side='left')
        end_idx = np.searchsorted(sub, end_date, side='right')
        if start_idx < end_idx:
            start_indices.append(start_idx)
            end_indices.append(end_idx)
    dfs = []
    for i in range(len(start_indices)):
        sub = pd.read_csv(file_paths[i])
        sub_range = (sub['Date'] >= start_date) & (sub['Date'] <= end_date)
        sub = sub[sub_range]
        dfs.append(sub)
    result = pd.concat(dfs)
    return result.sort_values(by='Date')

如上所述，我们将pandas的Series对象转换为Numpy数组，然后使用Numpy的searchsorted函数在每个块中查找开始和结束时间的索引，并将它们存储在start_indices和end_indices列表中。接下来，我们使用起始和结束索引从每个块中提取出数据。最后，我们将它们连接起来，并按照日期排序。

现在我们可以进行速度的比较。

start_date = pd.to_datetime('20210110')
end_date = pd.to_datetime('20210630')

%timeit query_data(start_date, end_date)
%timeit query_data_np(start_date, end_date)

输出：

100 loops, best of 5: 3.25 ms per loop
100 loops, best of 5: 2.29 ms per loop

使用Numpy来优化基于范围的合并确实提高了查询速度。

总结

通过使用Numpy数组和搜索方法，我们可以将基于范围的合并速度大大优化。使用pandas数据框和Numpy库可以使我们更轻松地处理大量数据，并使数据分析任务更容易实现。通过在pandas数据框和Numpy数组之间进行选择，我们可以更好地管理数据，使任务更加高效。