Режим строки как новый столбец в PySpark DataFrame - PullRequest
Новая
       14

Режим строки как новый столбец в PySpark DataFrame

1 голос
/ 04 июня 2019

Можно ли добавить новый столбец на основе максимума предыдущих столбцов, где предыдущие столбцы являются строковыми литералами.Рассмотрим следующий фрейм данных: 

df = spark.createDataFrame(
    [
        ('1',25000,"black","black","white"),
        ('2',16000,"red","black","white"),
    ],
    ['ID','cash','colour_body','colour_head','colour_foot']
)

Тогда целевой кадр должен выглядеть следующим образом: 

df = spark.createDataFrame(
    [
        ('1',25000,"black","black","white", "black" ),
        ('2',16000,"red","black","white", "white" ),
    ],
    ['ID','cash','colour_body','colour_head','colour_foot', 'max_v']
)

Если максимальный уровень обнаружения не обнаружен, следует использовать последний действительный цвет.

Есть ли какая-нибудь возможность счетчика или udf?

Ответы [ 2 ]

1 голос
/ 06 июня 2019

Для общего случая любого числа столбцов решение udf с помощью @ cs95 является подходящим вариантом.

Однако в этом конкретном случае, когда у вас есть только 3 столбца, вы можете упростить логику, используя pyspark.sql.functions.when, что будет более эффективным, чем udf. 

from pyspark.sql.functions import col, when

def mode_of_3_cols(body, head, foot):
    return(
        when(
            (body == head)|(body == foot), 
            body
        ).when(
            (head == foot),
            head
        ).otherwise(foot)
    )

df.withColumn(
    "max_v", 
    mode_of_3_cols(col("colour_body"), col("colour_head"), col("colour_foot"))
).show()
#+---+-----+-----------+-----------+-----------+-----+
#| ID| cash|colour_body|colour_head|colour_foot|max_v|
#+---+-----+-----------+-----------+-----------+-----+
#|  1|25000|      black|      black|      white|black|
#|  2|16000|        red|      black|      white|white|
#+---+-----+-----------+-----------+-----------+-----+

Вам просто нужно проверить, равны ли какие-либо два столбца - если да, то это значение должно быть режимом. Если нет, верните последний столбец.

1 голос
/ 04 июня 2019

Определите UDF около statistics.mode для вычисления построчного режима с требуемой семантикой: 

import statistics

from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

def mode(*x):
    try:
        return statistics.mode(x)
    except statistics.StatisticsError:
        return x[-1]

mode = udf(mode, StringType())


df.withColumn("max_v", mode(*[col(c) for c in df.columns if 'colour' in c])).show()

+---+-----+-----------+-----------+-----------+-----+
| ID| cash|colour_body|colour_head|colour_foot|max_v|
+---+-----+-----------+-----------+-----------+-----+
|  1|25000|      black|      black|      white|black|
|  2|16000|        red|      black|      white|white|
+---+-----+-----------+-----------+-----------+-----+
...