代码之家 › 专栏 › 技术社区 › Demetri Pananos

从列表列提取元素

pyspark

0

Demetri Pananos · 技术社区 · 8 年前

ix users .我想用 用户

 l = [('store1', [[0, 1], [2, 3]], [1001, 1002, 1003, 1004, 1005, 1006])]
 df = sqlContext.createDataFrame(l,('store','ix','users'))

+------+--------------------+--------------------+
| store|                  ix|               users|
+------+--------------------+--------------------+
|store1|[WrappedArray(0, ...|[1001, 1002, 1003...|
+------+--------------------+--------------------+

我希望输出

want =  [('store1', [[0, 1], [2, 3]], [1001, 1002, 1003, 1004, 1005, 1006], [[1001,1002],[1003,1004]])]
want = sqlContext.createDataFrame(want,('store','ix','users','new_col'))

如果我做的话 users[ix]

最好的方法是什么?

1 回复 | 直到 8 年前

1

Pierre Gourseaud 8 年前

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, IntegerType

def get_users_from_ix(ixs_array, users):
    output_array = []
    for ixs in ixs_array:
        output = []
        for ix in ixs:
            output.append(users[ix])
        output_array.append(output)
    return output_array

get_users_from_idx_udf = F.udf(get_users_from_idx, ArrayType(ArrayType(IntegerType())))

df_result = df.select('*', get_users_from_idx_udf(F.col('ix'), F.col('users')))

这应该有效。