如何花式索引到 PySpark DataFrame 中？

Question

我有一个像这样的 PySpark DataFrame：

|id1 | id2 |
------------
|100 |abcd |
|201 |efgh |
|112 |ijkl |
|133 |mnop |

还有 python 个 id2 列表，例如：

mylist = ['abcd', 'ijkl', 'efgh', 'efgh', 'abcd', 'ijkl', 'ijkl']

我想使用 pyspark 数据框将列表中的值映射到相应的 id1 值，因此：

desired_result = [100, 112, 201, 201, 100, 112, 112]

我现在正在使用 Pandas 进行此操作，但这种方式需要很长时间，因为查找 table 有数百万行。我如何在 PySpark 中执行此操作？

Answer 1

从您的列表中创建一个 DataFrame，df2 下面：

a = [(100, 'abcd'),(201, 'efgh'),(112, 'ijkl'),(133, 'mnop'), (245, 'qrt')]
df1 = spark.createDataFrame(a, ['id1', 'id2'])
mylist = ['abcd', 'ijkl', 'efgh', 'efgh', 'abcd', 'ijkl', 'ijkl']
ml = [(idx, k) for idx, k in enumerate(mylist)]
print(ml)
df1.show()
df2 = spark.createDataFrame(ml, ['id', 'name'])
df2.show()

它们如下所示：

[(0, 'abcd'), (1, 'ijkl'), (2, 'efgh'), (3, 'efgh'), (4, 'abcd'), (5, 'ijkl'), (6, 'ijkl')]
+---+----+
|id1| id2|
+---+----+
|100|abcd|
|201|efgh|
|112|ijkl|
|133|mnop|
|245| qrt|
+---+----+

+---+----+
| id|name|
+---+----+
|  0|abcd|
|  1|ijkl|
|  2|efgh|
|  3|efgh|
|  4|abcd|
|  5|ijkl|
|  6|ijkl|

加入DataFrames并使用id列保持原来的顺序

df = df2.join(df1, df1.id2==df2.name).orderBy('id')
df.show()
[r.id1 for r in df.select('id1').collect()]

最终输出：

+---+----+---+----+
| id|name|id1| id2|
+---+----+---+----+
|  0|abcd|100|abcd|
|  1|ijkl|112|ijkl|
|  2|efgh|201|efgh|
|  3|efgh|201|efgh|
|  4|abcd|100|abcd|
|  5|ijkl|112|ijkl|
|  6|ijkl|112|ijkl|
+---+----+---+----+

[100, 112, 201, 201, 100, 112, 112]

如何花式索引到 PySpark DataFrame 中？

How to Fancy-Index Into a PySpark DataFrame?

python

indexing

pyspark

databricks