Возвращать сложные типы из PySpark - PullRequest
Новая
       32

Возвращать сложные типы из PySpark

0 голосов
/ 28 февраля 2019

Я пытаюсь вернуть матрицу значений из UDF Pandas Grouped Map.

При определении схемы я даю Array (Array (DoubleType ())), но это вызывает возвраттолько первая строка матрицы со всеми другими значениями как None.

Вот что я попробовал: 

myschema = StructType([
  StructField('my_id',  StringType()),
  StructField('matrix', ArrayType(ArrayType(DoubleType())) )
]);
cols = list(map(lambda s: s.__dict__['name'], myschema))

@F.pandas_udf(myschema, F.PandasUDFType.GROUPED_MAP)
def my_function(data):
  myID = 'Hello'

  matrix = [
    [1.1, 2.2, 3.3],
    [4.4, 5.5, 6.6],
    [7.7, 8.8, 9.9]
  ]

  return pd.DataFrame([[ myID, matrix ]], columns=cols)

df = spark.createDataFrame(pd.DataFrame(['id1', 'id2'], columns=['ID']))
df.groupBy('ID').apply(my_function).collect()

И в результате у меня есть:

Row (my_id = u'Hello ', matrix = [[1.1, 2.2, 3.3], None, None])

Конечно, возврат не вложенного массива работает без проблем. 

myschema = StructType([
  StructField('my_id',  StringType()),
  StructField('matrix', ArrayType(DoubleType()))
]);
cols = list(map(lambda s: s.__dict__['name'], myschema))

@F.pandas_udf(myschema, F.PandasUDFType.GROUPED_MAP)
def my_function(data):
  myID = 'Hello'

  matrix = [1.1, 2.2, 3.3]

  return pd.DataFrame([[ myID, matrix ]], columns=cols)

df = spark.createDataFrame(pd.DataFrame(['id1', 'id2'], columns=['ID']))
df.groupBy('ID').apply(my_function).collect()

И результат:

Строка (my_id = u'Hello ', matrix = = [1.1, 2.2, 3.3])

Есть личто-то не так со схемой?

...