Question

我正在尝试从分组地图熊猫UDF返回值矩阵。

定义架构时，我要提供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'，矩阵= [1.1，2.2，3.3]）

架构有问题吗？

从PySpark返回复杂类型

0 个答案: