Как вернуть массив структуры или класса из UDF в значение столбца dataframe? - PullRequest
Новая
       30

Как вернуть массив структуры или класса из UDF в значение столбца dataframe?

0 голосов
/ 14 ноября 2018 
d = [{'ID': '1', 'pID': 1000, 'startTime':'2018.07.02T03:34:20', 'endTime':'2018.07.03T02:40:20'}, {'ID': '1', 'pID': 1000, 'startTime':'2018.07.02T03:45:20', 'endTime':'2018.07.03T02:50:20'}, {'ID': '2', 'pID': 2000, 'startTime':'2018.07.02T03:34:20', 'endTime':'2018.07.03T02:40:20'}, {'ID': '2', 'pID': 2000, 'startTime':'2018.07.02T03:45:20', 'endTime':'2018.07.03T02:50:20'}]

df = spark.createDataFrame(d)

Dates = namedtuple("Dates", "startTime endTime")


def MergeAdjacentUsage(timeSets):
  DatesArray = []
  for times in timeSets:
    DatesArray.append(Dates(startTime=times.startTime, endTime=times.endTime))
  return DatesArray


MergeAdjacentUsages = udf(MergeAdjacentUsage,ArrayType(Dates()))

df1=df.groupBy(['ID','pID']).agg(MergeAdjacentUsages(F.collect_list(struct('startTime','endTime'))).alias("Times"))

display(df1)

Все, что я хочу, это установить значение столбца в массив stuct, который возвращается UDF.Это дает мне ошибку как:

TypeError: new () принимает ровно 3 аргумента (1 дано)

TypeError Traceback (последний вызов последний) в() 22 return DatesArray 23 ---> 24 MergeAdjacentUsages = udf (MergeAdjacentUsage, ArrayType (Dates ())) 25 26 df1 = df.groupBy (['ID', 'pID']). Agg (MergeAdjacentUsages (F.collect_list)(struct ('startTime', 'endTime'))). alias ("Times"))

Любая помощь, идея или подсказка будут оценены.

1 Ответ

0 голосов
/ 17 ноября 2018

pyspark не допускает определяемые пользователем объекты Class как типы столбцов Dataframe.Вместо этого нам нужно создать StructType, который можно использовать аналогично классу / именованному кортежу в python.

Например: 

from pyspark.sql.types import *
from pyspark.sql.functions import udf
from pyspark.sql import functions as F
# from pyspark.sql.functions import *

d = [{'ID': '1', 'pID': 1000, 'startTime': '2018.07.02T03:34:20', 'endTime': '2018.07.03T02:40:20'},
     {'ID': '1', 'pID': 1000, 'startTime': '2018.07.02T03:45:20', 'endTime': '2018.07.03T02:50:20'},
     {'ID': '2', 'pID': 2000, 'startTime': '2018.07.02T03:34:20', 'endTime': '2018.07.03T02:40:20'},
     {'ID': '2', 'pID': 2000, 'startTime': '2018.07.02T03:45:20', 'endTime': '2018.07.03T02:50:20'}]

df = spark.createDataFrame(d)

# Dates = namedtuple("Dates", "startTime endTime")

schema = ArrayType(StructType([
    StructField("startTime", StringType(), False),
    StructField("endTime", StringType(), False)
]))


MergeAdjacentUsages = udf(lambda xs: xs, schema)

df1 = df.groupBy(['ID', 'pID']).agg(MergeAdjacentUsages(
    F.collect_list(F.struct('startTime', 'endTime'))).alias("Times"))
df1.show(truncate=False)

+---+----+----------------------------------------------------------------------------------------+
|ID |pID |Times                                                                                   |
+---+----+----------------------------------------------------------------------------------------+
|2  |2000|[[2018.07.02T03:34:20, 2018.07.03T02:40:20], [2018.07.02T03:45:20, 2018.07.03T02:50:20]]|
|1  |1000|[[2018.07.02T03:34:20, 2018.07.03T02:40:20], [2018.07.02T03:45:20, 2018.07.03T02:50:20]]|
+---+----+----------------------------------------------------------------------------------------+

Надеюсь, это поможет!

...