从列表列

Question

我的列是列表ix，另一列是列表users。我想将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]，那几乎就是这样。

实现这一目标的最佳方法是什么？

Answer 1

我认为你必须用udf来做。

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')))

这应该有用。