并行計算框架Polars、Dask的數據處理性能對比

作者：Luís Oliveira 2023-07-10 13:51:45

本文我們使用兩個類似的腳本來執行提取、轉換和加載(ETL)過程。

在Pandas 2.0發布以后，我們發布過一些評測的文章，這次我們看看，除了Pandas以外，常用的兩個都是為了大數據處理的并行數據框架的對比測試。

本文我們使用兩個類似的腳本來執行提取、轉換和加載(ETL)過程。

測試內容

這兩個腳本主要功能包括:

從兩個parquet 文件中提取數據，對于小型數據集，變量path1將為“yellow_tripdata/ yellow_tripdata_2014-01”，對于中等大小的數據集，變量path1將是“yellow_tripdata/yellow_tripdata”。對于大數據集，變量path1將是“yellow_tripdata/yellow_tripdata*.parquet”;

進行數據轉換:a)連接兩個DF，b)根據PULocationID計算行程距離的平均值，c)只選擇某些條件的行，d)將步驟b的值四舍五入為2位小數，e)將列“trip_distance”重命名為“mean_trip_distance”，f)對列“mean_trip_distance”進行排序。

將最終的結果保存到新的文件。

腳本

1、Polars

數據加載讀取

def extraction():
    """
    Extract two datasets from parquet files
    """
    path1="yellow_tripdata/yellow_tripdata_2014-01.parquet"
    df_trips= pl_read_parquet(path1,)
    path2 = "taxi+_zone_lookup.parquet"
    df_zone = pl_read_parquet(path2,)
 
    return df_trips, df_zone
 
 def pl_read_parquet(path, ):
    """
    Converting parquet file into Polars dataframe
    """
    df= pl.scan_parquet(path,)
    return df

轉換函數

def transformation(df_trips, df_zone):
    """
    Proceed to several transformations
    """
    df_trips= mean_test_speed_pl(df_trips, )
     
    df = df_trips.join(df_zone,how="inner", left_on="PULocationID", right_on="LocationID",)
    df = df.select(["Borough","Zone","trip_distance",])
   
    df = get_Queens_test_speed_pd(df)
    df = round_column(df, "trip_distance",2)
    df = rename_column(df, "trip_distance","mean_trip_distance")
 
    df = sort_by_columns_desc(df, "mean_trip_distance")
    return df
 
 
 def mean_test_speed_pl(df_pl,):
    """
    Getting Mean per PULocationID
    """
    df_pl = df_pl.groupby('PULocationID').agg(pl.col(["trip_distance",]).mean())
    return df_pl
 
 def get_Queens_test_speed_pd(df_pl):
    """
    Only getting Borough in Queens
    """
 
    df_pl = df_pl.filter(pl.col("Borough")=='Queens')
 
    return df_pl
 
 def round_column(df, column,to_round):
    """
    Round numbers on columns
    """
    df = df.with_columns(pl.col(column).round(to_round))
    return df
 
 def rename_column(df, column_old, column_new):
    """
    Renaming columns
    """
    df = df.rename({column_old: column_new})
    return df
 
 def sort_by_columns_desc(df, column):
    """
    Sort by column
    """
    df = df.sort(column, descending=True)
    return df

保存

def loading_into_parquet(df_pl):
    """
    Save dataframe in parquet
    """
    df_pl.collect(streaming=True).write_parquet(f'yellow_tripdata_pl.parquet')

其他代碼

import polars as pl
 import time
 
 def pl_read_parquet(path, ):
    """
    Converting parquet file into Polars dataframe
    """
    df= pl.scan_parquet(path,)
    return df
 
 def mean_test_speed_pl(df_pl,):
    """
    Getting Mean per PULocationID
    """
    df_pl = df_pl.groupby('PULocationID').agg(pl.col(["trip_distance",]).mean())
    return df_pl
 
 def get_Queens_test_speed_pd(df_pl):
    """
    Only getting Borough in Queens
    """
 
    df_pl = df_pl.filter(pl.col("Borough")=='Queens')
 
    return df_pl
 
 def round_column(df, column,to_round):
    """
    Round numbers on columns
    """
    df = df.with_columns(pl.col(column).round(to_round))
    return df
 
 def rename_column(df, column_old, column_new):
    """
    Renaming columns
    """
    df = df.rename({column_old: column_new})
    return df
 
 
 def sort_by_columns_desc(df, column):
    """
    Sort by column
    """
    df = df.sort(column, descending=True)
    return df
 
 
 def main():
     
    print(f'Starting ETL for Polars')
    start_time = time.perf_counter()
 
    print('Extracting...')
    df_trips, df_zone =extraction()
        
    end_extract=time.perf_counter() 
    time_extract =end_extract- start_time
 
    print(f'Extraction Parquet end in {round(time_extract,5)} seconds')
    print('Transforming...')
    df = transformation(df_trips, df_zone)
    end_transform = time.perf_counter() 
    time_transformation =time.perf_counter() - end_extract
    print(f'Transformation end in {round(time_transformation,5)} seconds')
    print('Loading...')
    loading_into_parquet(df,)
    load_transformation =time.perf_counter() - end_transform
    print(f'Loading end in {round(load_transformation,5)} seconds')
    print(f"End ETL for Polars in {str(time.perf_counter()-start_time)}")
 
 
 if __name__ == "__main__":
     
    main()

2、Dask

函數功能與上面一樣，所以我們把代碼整合在一起：

import dask.dataframe as dd
 from dask.distributed import Client
 import time
 
 def extraction():
    path1 = "yellow_tripdata/yellow_tripdata_2014-01.parquet"
    df_trips = dd.read_parquet(path1)
    path2 = "taxi+_zone_lookup.parquet"
    df_zone = dd.read_parquet(path2)
 
    return df_trips, df_zone
 
 def transformation(df_trips, df_zone):
    df_trips = mean_test_speed_dask(df_trips)
    df = df_trips.merge(df_zone, how="inner", left_on="PULocationID", right_on="LocationID")
    df = df[["Borough", "Zone", "trip_distance"]]
 
    df = get_Queens_test_speed_dask(df)
    df = round_column(df, "trip_distance", 2)
    df = rename_column(df, "trip_distance", "mean_trip_distance")
 
    df = sort_by_columns_desc(df, "mean_trip_distance")
    return df
 
 def loading_into_parquet(df_dask):
    df_dask.to_parquet("yellow_tripdata_dask.parquet", engine="fastparquet")
 
 def mean_test_speed_dask(df_dask):
    df_dask = df_dask.groupby("PULocationID").agg({"trip_distance": "mean"})
    return df_dask
 
 def get_Queens_test_speed_dask(df_dask):
    df_dask = df_dask[df_dask["Borough"] == "Queens"]
    return df_dask
 
 def round_column(df, column, to_round):
    df[column] = df[column].round(to_round)
    return df
 
 def rename_column(df, column_old, column_new):
    df = df.rename(columns={column_old: column_new})
    return df
 
 def sort_by_columns_desc(df, column):
    df = df.sort_values(column, ascending=False)
    return df
 
 
 
 def main():
    print("Starting ETL for Dask")
    start_time = time.perf_counter()
 
    client = Client() # Start Dask Client
 
    df_trips, df_zone = extraction()
 
    end_extract = time.perf_counter()
    time_extract = end_extract - start_time
 
    print(f"Extraction Parquet end in {round(time_extract, 5)} seconds")
    print("Transforming...")
    df = transformation(df_trips, df_zone)
    end_transform = time.perf_counter()
    time_transformation = time.perf_counter() - end_extract
    print(f"Transformation end in {round(time_transformation, 5)} seconds")
    print("Loading...")
    loading_into_parquet(df)
    load_transformation = time.perf_counter() - end_transform
    print(f"Loading end in {round(load_transformation, 5)} seconds")
    print(f"End ETL for Dask in {str(time.perf_counter() - start_time)}")
 
    client.close() # Close Dask Client
 
 if __name__ == "__main__":
    main()