Использовать результат искровой функции как вход другой функции

Question

В моем приложении Spark у меня есть фрейм данных с такой информацией, как

+------------------+---------------+
|      labels      | labels_values |
+------------------+---------------+
| ['l1','l2','l3'] |           000 |
| ['l3','l4','l5'] |           100 |
+------------------+---------------+

. Я пытаюсь создать с заданным именем метки в качестве входного значения столбец single_label_value, который принимает значение для этого метка из столбца labels_values.

Например, для label='l3' я хотел бы получить этот вывод:

+------------------+---------------+--------------------+
|      labels      | labels_values | single_label_value |
+------------------+---------------+--------------------+
| ['l1','l2','l3'] |           000 |                  0 |
| ['l3','l4','l5'] |           100 |                  1 |
+------------------+---------------+--------------------+

Вот что я пытаюсь использовать:

selected_label='l3'
label_position = F.array_position(my_df.labels, selected_label)
my_df= my_df.withColumn(
    "single_label_value", 
    F.substring(my_df.labels_values, label_position, 1)
)

Но я получаю ошибку, потому что функция подстроки не любит аргумент label_position.

Есть ли способ объединить эти выходные данные функции без записи udf?

Answer 1

Надеюсь, это сработает для вас.

from pyspark.sql import SparkSession
from pyspark.sql.functions import *

spark=SparkSession.builder.getOrCreate()
mydata=[[['l1','l2','l3'],'000'], [['l3','l4','l5'],'100']]

df = spark.createDataFrame(mydata,schema=["lebels","lebel_values"]) 

selected_label='l3'

df2=df.select(
        "*",
        (array_position(df.lebels,selected_label)-1).alias("pos_val"))

df2.createOrReplaceTempView("temp_table")

df3=spark.sql("select *,substring(lebel_values,pos_val,1) as val_pos from temp_table")

df3.show()


+------------+------------+-------+-------+
|      lebels|lebel_values|pos_val|val_pos|
+------------+------------+-------+-------+
|[l1, l2, l3]|         000|      2|      0|
|[l3, l4, l5]|         100|      0|      1|
+------------+------------+-------+-------+

Это дает местоположение значения. Если вам нужен точный индекс, вы можете использовать -1 из этого значения.

- Отредактированный anser -> Работал с временным представлением. Все еще ищу решение, используя опцию withColumn. Я надеюсь, это поможет вам сейчас. Edit2 -> Ответить, используя фрейм данных.

df2=df.select(
        "*",
        (array_position(df.lebels,selected_label)-1).astype("int").alias("pos_val")
        )

df3=df2.withColumn("asked_col",expr("substring(lebel_values,pos_val,1)"))
df3.show()

Answer 2

Попробуйте, может быть:

import pyspark.sql.functions as f
from pyspark.sql.functions import *

selected_label='l3'
df=df.withColumn('single_label_value', f.substring(f.col('labels_values'), array_position(f.col('labels'), lit(selected_label))-1, 1))

df.show()

(для версии с искрой> = 2.4)

Я думаю, lit() была функция, которую вы пропустили - вы можете использовать ее для передачи константных значений в spark датафреймы.