Более эффективный способ прокрутки PySpark DataFrame и создания новых столбцов

Question

Я конвертирую код, написанный с помощью Pandas в PySpark. В коде есть много циклов for, чтобы создать переменное количество столбцов в зависимости от заданных пользователем входов.

Я использую Спарк 1.6.x, с помощью следующего кода образца:

from pyspark.sql import SQLContext 
from pyspark.sql import functions as F 
import pandas as pd 
import numpy as np 

# create a Pandas DataFrame, then convert to Spark DataFrame 
test = sqlContext.createDataFrame(pd.DataFrame({'val1': np.arange(1,11)})) 

Который оставляет меня

+----+ 
|val1| 
+----+ 
| 1| 
| 2| 
| 3| 
| 4| 
| 5| 
| 6| 
| 7| 
| 8| 
| 9| 
| 10| 
+----+ 

I цикла много в коде, например, внизу:

for i in np.arange(2,6).tolist(): 
    test = test.withColumn('val_' + str(i), F.lit(i ** 2) + test.val1) 

Какие результаты в:

+----+-----+-----+-----+-----+ 
|val1|val_2|val_3|val_4|val_5| 
+----+-----+-----+-----+-----+ 
| 1| 5| 10| 17| 26| 
| 2| 6| 11| 18| 27| 
| 3| 7| 12| 19| 28| 
| 4| 8| 13| 20| 29| 
| 5| 9| 14| 21| 30| 
| 6| 10| 15| 22| 31| 
| 7| 11| 16| 23| 32| 
| 8| 12| 17| 24| 33| 
| 9| 13| 18| 25| 34| 
| 10| 14| 19| 26| 35| 
+----+-----+-----+-----+-----+ 

** Вопрос: ** Как я могу переписать вышеуказанный цикл, чтобы быть более эффективным?

Я заметил, что мой код работает медленнее, поскольку Spark проводит много времени на каждой группе циклов (даже на небольших наборах данных, таких как 2 ГБ ввода текста).

Благодаря

Answer 1

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

df = spark.range(1, 11).toDF("val1") 

def make_col(i): 
    return (F.pow(F.lit(i), 2) + F.col("val1")).alias("val_{0}".format(i)) 

spark.range(1, 11).toDF("val1").select("*", *(make_col(i) for i in range(2, 6))) 

Я бы также избегал использования типов NumPy. Инициализация объектов NumPy обычно более дорога по сравнению с обычными объектами Python, а Spark SQL не поддерживает типы NumPy, поэтому требуются дополнительные преобразования.

Answer 2

Один withColumn будет работать на всей РДУ. Поэтому обычно не рекомендуется использовать метод для каждого столбца, который вы хотите добавить. Существует способ, которым вы работаете со столбцами и их данными внутри функции карты. Поскольку здесь используется одна функция карты, код для добавления нового столбца и его данных будет выполняться параллельно.

a. вы можете собирать новые значения на основе расчетов

b. Добавьте эти новые значения столбцов в основной RDD, как показано ниже

val newColumns: Seq[Any] = Seq(newcol1,newcol2) 
Row.fromSeq(row.toSeq.init ++ newColumns) 

Вот ряд, является ссылкой на строку в методе карты

гр. Создать новую схему, указанную ниже

val newColumnsStructType = StructType{Seq(new StructField("newcolName1",IntegerType),new StructField("newColName2", IntegerType)) 

d. Добавить в старую схему

val newSchema = StructType(mainDataFrame.schema.init ++ newColumnsStructType) 

e. Создать новую dataframe с новыми столбцами

val newDataFrame = sqlContext.createDataFrame(newRDD, newSchema)