Подстрока Pyspark одного столбца на основе длины другого столбца - PullRequest
Новая
       35

Подстрока Pyspark одного столбца на основе длины другого столбца

0 голосов
/ 25 февраля 2019

Использование Pyspark 2.2

У меня есть спарк DataFrame с несколькими столбцами.Мне нужно ввести 2 столбца в UDF и вернуть 3-й столбец

Ввод: 

+-----+------+
|col_A| col_B|
+-----+------+
|  abc|abcdef|
|  abc|     a|
+-----+------+

И col_A, и col_B StringType()

Желаемыйвывод: 

+-----+------+-------+
|col_A| col_B|new_col|
+-----+------+-------+
|  abc|abcdef|    abc|
|  abc|     a|      a|
+-----+------+-------+

Я хочу, чтобы new_col была подстрокой col_A с длиной col_B.

Я пытался 

udf_substring = F.udf(lambda x: F.substring(x[0],0,F.length(x[1])), StringType())
df.withColumn('new_col', udf_substring([F.col('col_A'),F.col('col_B')])).show()

Ноэто дает TypeError: Column is not iterable.

Есть идеи, как сделать такие манипуляции?

1 Ответ

0 голосов
/ 25 февраля 2019

Здесь две главные вещи не так.

  • Во-первых, вы определили, что ваш udf принимает один входной параметр, когда он должен принимать значение 2.
  • Во-вторых, вы не можете использовать функции API в udf,(Вызов udf сериализуется в python, поэтому вам нужно использовать синтаксис и функции python.)

Вот правильная реализация udf для этой проблемы: 

import pyspark.sql.functions as F

def my_substring(a, b):
    # You should add in your own error checking
    return a[:len(b)]

udf_substring = F.udf(lambda x, y: my_substring(a, b), StringType())

Изатем вызовите его, передав в качестве аргументов два столбца: 

df.withColumn('new_col', udf_substring(F.col('col_A'),F.col('col_B')))

Однако в этом случае вы можете сделать это без udf, используя метод, описанный в этого поста . 

df.withColumn(
    'new_col', 
    F.expr("substring(col_A,0,length(col_B))")
)
...