PySpark: получить n-е значение кортежа для UDF - PullRequest
Новая
       87

PySpark: получить n-е значение кортежа для UDF

0 голосов
/ 27 мая 2020

У меня есть функция UDF с выводом в формате кортежа. Я хочу применить этот UDF к моему входному столбцу, и в зависимости от того, что мне нужно, я хочу выбрать значение out1 или out2 в качестве значения для моего столбца.

Примерно так: 

def my_f(inp):
   return out1,out2

df =  df.withColumn('first_val', F.udf(my_f, StringType())(F.col('inp_col'))[0]
df =  df.withColumn('second_val', F.udf(my_f, StringType())(F.col('inp_col'))[1]

Я хочу, чтобы first_val col содержал первый элемент кортежа, а second_val - второй элемент кортежа. Этот код, конечно, не работает.

Я попробовал его, передав требуемую часть кортежа в качестве ввода функции, и он сработал. Примерно так: 

def my_f(inp, out='full'):
   if out=='first':
      return out1
   elif out=='second':
      return out2
   else: # case 'full'
   return out1,out2

df =  df.withColumn('first_val',  F.udf(my_f, StringType())(F.col('inp_col'), F.col('inp_col'))
df =  df.withColumn('second_val', F.udf(my_f, StringType())(F.col('inp_col'),'F.col('second'))

Но есть ли более простой способ получить n-й элемент кортежа в строке без передачи этого параметра?

1 Ответ

1 голос
/ 27 мая 2020

Если ваш UDF возвращает кортеж, вы должны изменить тип возвращаемого значения на ArrayType(StringType), предполагая, что вы возвращаете кортеж строк. Тогда вы сможете получить доступ к первому и второму элементу вашего кортежа, используя нотацию [n]. Вот пример: 

import pyspark.sql.functions as F
import pyspark.sql.types as T
...
@F.udf(T.ArrayType(T.StringType())
def my_f(inp):
   ...
   return (out1, out2)

df =  df.withColumn('first_val', my_f('inp_col')[0])
df =  df.withColumn('second_val', my_f('inp_col')[1])

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

import pyspark.sql.functions as F
import pyspark.sql.types as T
...
@F.udf(T.StructType([
    T.StructField("first", T.StringType()),
    T.StructField("second", T.IntegerType())
]))
def my_f(inp):
   ...
   return {"first": out1, "second": out2}

df =  df.withColumn('first_val', my_f('inp_col')["first"])
df =  df.withColumn('second_val', my_f('inp_col')["second"])
...