将标签字符串转换为二进制向量 pyspark

Question

I have data which looks like this:

| Id | ----Tags---- | some_text |
| 0  | <a><b>       | ex1       |
| 1  | <a><c>       | ex2       |
| 2  | <b><c>       | ex3       |

and I want it to end up looking like this:

| Id | a | b | c | some_text |
| 0  | 1 | 1 | 0 | ex1       |
| 1  | 1 | 0 | 1 | ex2       |
| 2  | 0 | 1 | 1 | ex3       |

and I would like to use pyspark for the solution. Any ideas on how to approach this problem?

Answer 1

If you don't already know the expected categorical values, you can use pyspark.sql.functions.udf to split and tags into an array of values and pyspark.sql.functions.explode function to convert them to columns. Then, you can pivot values to columns:

# required imports
import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType, StringType
import re

# regex pattern to split 'tagged values'
pat = re.compile('<(.*?)>')

#udf to split string to array of values
split_f = f.udf(lambda s: pat.split(s), ArrayType(StringType()))

# sample data
df = spark.createDataFrame([(0,'<a><b>','ex1'),(1,'<a><c>','ex2')], ['Id', '---Tags---', 'some_text'])

+---+----------+---------+
| Id|---Tags---|some_text|
+---+----------+---------+
|  0|    <a><b>|      ex1|
|  1|    <a><c>|      ex2|
+---+----------+---------+

df.withColumn('exploded', 
   F.explode(split_f(F.col('---Tags---'))))
  .groupby('Id').pivot('exploded').count().na.fill(0).show()

+---+---+---+---+
| Id|  a|  b|  c|
+---+---+---+---+
|  0|  1|  1|  0|
|  1|  1|  0|  1|
+---+---+---+---+