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)

Измененный DataFrame:

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.На столбец Country_ID следует ссылаться с помощью modified_df.

, поэтому измените

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

на

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

, и он должен работать.Поскольку df не имеет столбца Country_ID.