PySpark 无法访问使用 StringIndexer 添加的列

Question

PySpark - v2.4.0

我尝试将String 列Country 转换为Interger 列Country_ID，结果看起来不错。但是当我尝试访问Country_ID 列时，我得到了AnalysisException。

下面是数据框：

+------+-------+
|UserId|Country|
+------+-------+
|     1| Africa|
|     2| Africa|
|     3|     UK|
|     4|  Japan|
|     5|     UK|
|     6|  Japan|
|     7|  China|
+------+-------+

代码如下：

from pyspark.ml.feature import StringIndexer
indexer = StringIndexer(inputCol='Country', outputCol='Country_ID')
modified_df = indexer.fit(df).transform(df)

修改后的数据框：

modified_df.select('*').show()

+------+-------+----------+
|UserId|Country|Country_ID|
+------+-------+----------+
|     1| Africa|       1.0|
|     2| Africa|       1.0|
|     3|     UK|       0.0|
|     4|  Japan|       2.0|
|     5|     UK|       0.0|
|     6|  Japan|       2.0|
|     7|  China|       3.0|
+------+-------+----------+

过滤查询：

modified_df.select('UserId').filter(df['Country_ID'] == 2).show()

以下是例外：

AnalysisException: u'Cannot resolve column name "Country_ID" among (UserId, Country);'

我可以将该列视为 DataFrame 的一部分

modified_df.columns

给予，

['UserId', 'Country', 'Country_ID']

如何使 DataFrame 过滤条件适用于Country_ID？

Answer 1

我认为问题在于您尝试使用错误的Dataframe 引用访问该列。应该使用modified_df 引用列Country_ID。

所以改变

modified_df.select('UserId').filter(df['Country_ID'] == 2).show()

到

modified_df.select('UserId').filter(modified_df['Country_ID'] == 2).show()

它应该可以工作。因为df 没有任何列Country_ID。