使用Dask从CSV文件中采样精确行数

Question

我想使用Dask创建n行的子样本。我尝试了两种方法：

1.使用frac：

import dask.dataframe as dd    
read_path = ["test_data\\small1.csv", "test_data\\small2.csv", "huge.csv"]
df = dd.read_csv(read_path)
df = df.sample(frac=0.0001)
df = df.compute()

它的工作速度足够快-从100mil数据集中选择10000，持续16秒。但是它不能保证确切的行数——因为使用了frac，所以它会被四舍五入。

2.使用for循环：

nrows = 10000
res_df = []
length = csv_loader.get_length()
total_len = sum(length)
start = perf_counter()
inds = random.sample(range(total_len), total_len - nrows - len(length))
min_bound = 0
relative_inds = []
for leng in length:
    relative_inds.append(
        sorted([i - min_bound for i in inds if min_bound <= i < min_bound + leng])
    )
    min_bound += leng
for ind, fil in enumerate(read_path):
    res_df.append(dd.read_csv(fil, skiprows=relative_inds[ind], sample=1000000))

在这里，我计算需要跳过的行的索引，然后使用skiprows从csv加载。这个方法非常慢，如果我需要从一些小csv读取0行，有时会崩溃。但它保证了精确的行数。

关于使用Dask获得精确的行数，是否有一些快速解决方案？

Answer 1

我找到了解决办法：

total_len = get_total_length() #compute len of all data in csvs
frac = nrows / total_len

while int(total_len * frac) != nrows:
    counter = 1
    frac = nrows / (total_len - counter)
    counter += 1

    res_df = dd.read_csv(read_path)
    res_df = res_df.sample(frac=0.0001)
    res_df = res_df.compute()

您可以查看如何有效地计算csv访问下一个链接中的行数。