Spark pandas_udf не быстрее - PullRequest
Новая
       19

Spark pandas_udf не быстрее

1 голос
/ 08 июля 2019

Я столкнулся с серьезным преобразованием данных. В двух словах, у меня есть столбцы данных, каждый из которых содержит строки, которые соответствуют некоторым порядковым номерам. Например, HIGH, MID и LOW. Моя цель - отобразить эти строки в целые числа, которые сохранят порядок. В этом случае LOW -> 0, MID -> 1 и HIGH -> 2.

Вот простая функция, генерирующая такие данные: 

def fresh_df(N=100000, seed=None):
    np.random.seed(seed)
    feat1 = np.random.choice(["HI", "LO", "MID"], size=N)
    feat2 = np.random.choice(["SMALL", "MEDIUM", "LARGE"], size=N)

    pdf = pd.DataFrame({
        "feat1": feat1,
        "feat2": feat2
    })
    return spark.createDataFrame(pdf)

Мой первый подход был: 

feat1_dict = {"HI": 1, "MID": 2, "LO": 3}
feat2_dict = {"SMALL": 0, "MEDIUM": 1, "LARGE": 2}

mappings = {
    "feat1": F.create_map([F.lit(x) for x in chain(*feat1_dict.items())]),
    "feat2": F.create_map([F.lit(x) for x in chain(*feat2_dict.items())])
}

for col in df.columns:
    col_map = mappings[col]
    df = df.withColumn(col+"_mapped", col_map[df[col]])

Это работает, как и ожидалось, но в действительности это происходит медленно, и я хотел оптимизировать процесс. Я прочитал о pandas_udf, и это дало мне надежду. Вот модифицированный код: 

feats_dict = {
    "feat1": feat1_dict,
    "feat2": feat2_dict
}

for col_name in df.columns:
    @F.pandas_udf('integer', F.PandasUDFType.SCALAR)
    def map_map(col):
        return col.map(feats_dict[col_name])
    df = df.withColumn(col_name + "_mapped", map_map(df[col_name]))

Увы! При сравнении этих двух версий не было никакого улучшения с точки зрения времени выполнения. Я сравнил их на локальном экземпляре Spark (с использованием докера) и на кластере EMR из 5 узлов (с настройками по умолчанию).

Я создал блокнот , где вы можете увидеть весь код. В общем, я использовал следующий импорт: 

import numpy as np
import pandas as pd

from itertools import chain
from pyspark.sql import functions as F

Чего мне не хватает? Почему этот процесс такой медленный и почему нет улучшения при использовании pandas_udf?

1 Ответ

1 голос
/ 08 июля 2019

Почему так медленно? Поскольку Spark работает в JVM, а pyspark - нет (потому что это процесс python), и для того, чтобы сделать его возможным, необходимо переместить все сериализацию и десериализацию данных в JVM.

Вы можете отобразить значения с помощью функций when и otherwise и избежать процесса сериализации и десериализации, увеличивая производительность. 

import numpy as np
import pandas as pd
import pyspark.sql.functions as f
from pyspark.shell import spark


def fresh_df(n=100000, seed=None):
    np.random.seed(seed)
    feat1 = np.random.choice(["HI", "LO", "MID"], size=n)
    feat2 = np.random.choice(["SMALL", "MEDIUM", "LARGE"], size=n)

    pdf = pd.DataFrame({
        "feat1": feat1,
        "feat2": feat2
    })
    return spark.createDataFrame(pdf)


df = fresh_df()
df = df.withColumn('feat1_mapped', f
                   .when(df.feat1 == f.lit('HI'), 1)
                   .otherwise(f.when(df.feat1 == f.lit('MID'), 2).otherwise(3)))

df = df.withColumn('feat2_mapped', f
                   .when(df.feat2 == f.lit('SMALL'), 0)
                   .otherwise(f.when(df.feat2 == f.lit('MEDIUM'), 1).otherwise(2)))
df.show(n=20)

выход 

+-----+------+------------+------------+
|feat1| feat2|feat1_mapped|feat2_mapped|
+-----+------+------------+------------+
|   LO| SMALL|           3|           0|
|   LO|MEDIUM|           3|           1|
|  MID|MEDIUM|           2|           1|
|  MID| SMALL|           2|           0|
|  MID| LARGE|           2|           2|
|  MID| SMALL|           2|           0|
|   LO| SMALL|           3|           0|
|  MID| LARGE|           2|           2|
|  MID| LARGE|           2|           2|
|  MID| SMALL|           2|           0|
|  MID|MEDIUM|           2|           1|
|   LO| LARGE|           3|           2|
|   HI|MEDIUM|           1|           1|
|   LO| SMALL|           3|           0|
|   HI|MEDIUM|           1|           1|
|  MID| SMALL|           2|           0|
|  MID|MEDIUM|           2|           1|
|   HI| SMALL|           1|           0|
|   HI| LARGE|           1|           2|
|  MID| LARGE|           2|           2|
+-----+------+------------+------------+
...