为什么在使用dask时,zarr的性能比镶木地板好得多



当我使用dask对zarr数据和镶木地板数据运行基本相同的计算时,基于zarr的计算速度明显更快。为什么?可能是因为我创建镶木地板文件时做错了什么吗?

我在jupyter笔记本中复制了假数据的问题(见下文(,以说明我看到的行为。如果有人能深入了解为什么基于zarr的计算比基于镶木地板的计算快几个数量级,我将不胜感激。

我在现实生活中使用的数据是地球科学模型数据。特定的数据参数并不重要,但每个参数都可以被认为是一个具有纬度、经度和时间维度的数组。

为了生成zarr文件,我只需写出参数的多维结构及其维度。

为了生成拼花地板,我首先将三维参数数组"展平"为一维数组,它成为我的数据框中的一列。然后添加纬度、经度和时间列,然后将数据框写成镶木地板。


此单元格具有代码其余部分所需的所有导入:

import pandas as pd
import numpy as np
import xarray as xr
import dask
import dask.array as da
import intake
from textwrap import dedent

这个单元生成伪造的数据文件,总大小超过3 GB:

def build_data(lat_resolution, lon_resolution, ntimes):
"""Build a fake geographical dataset with ntimes time steps and 
resolution lat_resolution x lon_resolution"""
lats = np.linspace(-90.0+lat_resolution/2,
90.0-lat_resolution/2,
np.round(180/lat_resolution))
lons = np.linspace(-180.0+lon_resolution/2,
180-lon_resolution/2,
np.round(360/lon_resolution))
times = np.arange(start=1,stop=ntimes+1)
data = np.random.randn(len(lats),len(lons),len(times))
return lats,lons,times,data
def create_zarr_from_data_set(lats,lons,times,data,zarr_dir):
"""Write zarr from a data set corresponding to the data passed in."""
dar = xr.DataArray(data,
dims=('lat','lon','time'),
coords={'lat':lats,'lon':lons,'time':times},
name="data")
ds = xr.Dataset({'data':dar,
'lat':('lat',lats),
'lon':('lon',lons),
'time':('time',times)})
ds.to_zarr(zarr_dir)
def create_parquet_from_data_frame(lats,lons,times,data,parquet_file):
"""Write a parquet file from a dataframe corresponding to the data passed in."""
total_points = len(lats)*len(lons)*len(times)
# Flatten the data array
data_flat = np.reshape(data,(total_points,1))
# use meshgrid to create the corresponding latitude, longitude, and time 
# columns
mesh = np.meshgrid(lats,lons,times,indexing='ij')
lats_flat = np.reshape(mesh[0],(total_points,1))
lons_flat = np.reshape(mesh[1],(total_points,1))
times_flat = np.reshape(mesh[2],(total_points,1))
df = pd.DataFrame(data = np.concatenate((lats_flat,
lons_flat,
times_flat, 
data_flat),axis=1), 
columns = ["lat","lon","time","data"])
df.to_parquet(parquet_file,engine="fastparquet")
def create_fake_data_files():
"""Create zarr and parquet files with fake data"""
zarr_dir = "zarr"
parquet_file = "data.parquet"
lats,lons,times,data = build_data(0.1,0.1,31)
create_zarr_from_data_set(lats,lons,times,data,zarr_dir)
create_parquet_from_data_frame(lats,lons,times,data,parquet_file)
with open("data_catalog.yaml",'w') as f:
catalog_str = dedent("""
sources:
zarr:
args:
urlpath: "./{}"
description: "data in zarr format"
driver: intake_xarray.xzarr.ZarrSource
metadata: {{}}
parquet:
args:
urlpath: "./{}"
description: "data in parquet format"
driver: parquet
""".format(zarr_dir,parquet_file))
f.write(catalog_str)

##
# Generate the fake data
##
create_fake_data_files()

我对parquet和zarr文件进行了几种不同的计算,但为了简单起见,在本例中,我只提取特定时间、纬度和经度的单个参数值。

该单元构建用于计算的zarr和parquet有向无环图(DAG(:

# pick some arbitrary point to pull out of the data
lat_value = -0.05
lon_value = 10.95
time_value = 5
# open the data
cat = intake.open_catalog("data_catalog.yaml")
data_zarr = cat.zarr.to_dask()
data_df = cat.parquet.to_dask()
# build the DAG for getting a single point out of the zarr data
time_subset = data_zarr.where(data_zarr.time==time_value,drop=True)
lat_condition = da.logical_and(time_subset.lat < lat_value + 1e-9, time_subset.lat > lat_value - 1e-9)
lon_condition = da.logical_and(time_subset.lon < lon_value + 1e-9, time_subset.lon > lon_value - 1e-9)
geo_condition = da.logical_and(lat_condition,lon_condition)
zarr_subset = time_subset.where(geo_condition,drop=True)
# build the DAG for getting a single point out of the parquet data
parquet_subset = data_df[(data_df.lat > lat_value - 1e-9) & 
(data_df.lat < lat_value + 1e-9) &
(data_df.lon > lon_value - 1e-9) & 
(data_df.lon < lon_value + 1e-9) &
(data_df.time == time_value)]

当我针对每个DAG的计算运行时间时,我得到的时间非常不同。基于zarr的子集只需不到一秒钟的时间。基于镶木地板的子集需要15-30秒。

该单元格进行基于zarr的计算:

%%time
zarr_point = zarr_subset.compute()

基于Zarr的计算时间:

CPU times: user 6.19 ms, sys: 5.49 ms, total: 11.7 ms
Wall time: 12.8 ms

此单元格进行基于镶木地板的计算:

%%time
parquet_point = parquet_subset.compute()

基于镶木地板的计算时间:

CPU times: user 18.2 s, sys: 28.1 s, total: 46.2 s
Wall time: 29.3 s

正如你所看到的,基于zarr的计算要快得多。为什么?

很高兴看到fastparquetzarrintake在同一个问题中使用!

TL;这里的DR是:使用适合您任务的正确数据模型。

此外,值得指出的是,zarr数据集为1.5GB,blosc/lz4压缩为512个块,而拼花地板数据集为1.8GB,快速压缩为5个块,其中压缩都是默认值。随机数据压缩不好,坐标压缩不好。

zarr是一种面向数组的格式,可以在任何维度上分块,这意味着,要读取一个点,只需要元数据(非常简短的文本(和包含它的一个块在这种情况下需要解压缩。数据块的索引是隐式的。

镶木地板是一种面向列的格式。为了找到一个特定的点,根据坐标列的组织方式,您可以忽略每个块的最小/最大列元数据中的一些块,然后为随机数据加载列块并解压缩。您需要自定义逻辑才能选择块同时加载到多个列上,而Dask目前没有实现这一点(如果不仔细重新排序数据,这是不可能的(。拼花地板的元数据比zarr大得多,但在这种情况下两者都无关紧要——如果你有很多变量或更多坐标,这可能会成为拼花地板额外的问题。

在这种情况下,zarr的随机访问速度会快得多,但读取所有数据并没有根本的不同,因为两者都必须将所有字节加载到磁盘上并解压缩到浮点中,而且在这两种情况下,坐标数据加载得很快。然而,未压缩数据帧的内存表示比未压缩数组大得多,因为现在每个坐标都有数组,而不是1D小数组,数组的点数与随机数据相同;此外,通过对小数组进行索引以在数组情况下获得正确的坐标,并与数据帧情况下每个点的每个lat/lon值进行比较,可以找到特定的点。

最新更新