Блог

У меня есть фрейм данных PySpark

 simpleData = [("person0",10, 10), 
    ("person1",1, 1), 
    ("person2",1, 0), 
    ("person3",5, 1), 
  ]
columns= ["persons_name","A", 'B']
exp = spark.createDataFrame(data = simpleData, schema = columns)

exp.printSchema()
exp.show()
 

Это выглядит так

 root
 |-- persons_name: string (nullable = true)
 |-- A: long (nullable = true)
 |-- B: long (nullable = true)
 |-- total: long (nullable = true)

 ------------ --- --- 
|persons_name|  A|  B|
 ------------ --- --- 
|     person0| 10| 10|
|     person1|  1|  1|   
|     person2|  1|  0|    
|     person3|  5|  1|    
 ------------ --- --- 
 

Теперь я хочу, чтобы к значениям столбцов A и B применялось пороговое значение 2, так что любое значение в столбце, меньшее порогового значения, становится 0, а значения, превышающие пороговое значение, становятся 1.

Конечный результат должен выглядеть примерно так-

  ------------ --- --- 
|persons_name|  A|  B|
 ------------ --- --- 
|     person0|  1|  1|
|     person1|  0|  0|   
|     person2|  0|  0|    
|     person3|  1|  0|    
 ------------ --- --- 
 

Как я могу этого добиться?

 threshold = 2
exp.select(
    [(F.col(col) > F.lit(threshold)).cast('int').alias(col) for col in ['A', 'B']]
)