PySpark transform dataframe

Question

Допустим, у меня есть следующие данные в кадре данных receipts:

Id | Фрукты

---

1 | ['яблоко', 'банан']

2 | [ 'Яблоко']

3 | [ 'Груша']

4 | ['груша', 'банан']

И я хочу превратить эти данные в матрицу, которая будет выглядеть следующим образом:

| A | Б | P |

---

| 1 | 1 | 0 |

| 1 | 0 | 0 |

| 0 | 0 | 1 |

| 0 | 1 | 1 |

(1, если присутствует фрукт, в противном случае 0)

Кто-нибудь знает эффективный способ сделать это? Моя настоящая проблема состоит из 60 000 "фруктов" и более 10 миллионов идентификаторов.

Answer 1

from pyspark.ml.feature import CountVectorizer
colorVectorizer = CountVectorizer(inputCol="Fruits",outputCol="Fruits_OneHotEncoded", vocabSize=60, minDF=1.0)
colorVectorizer_model = colorVectorizer.fit(df)
df_ohe = colorVectorizer_model.transform(df)
import numpy as np
x_3d = np.array(df_ohe.select('Fruits_OneHotEncoded').collect())
rows, idx, vocabsize = x_3d.shape
X = x_3d.reshape(rows, features)
X

Попробуйте приведенный выше код.