PySpark - список, созданный в столбце dataframe, имеет тип String вместо Integer

Question

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

values = [('A',8),('B',7)]
df = sqlContext.createDataFrame(values,['col1','col2'])
df.show()
+----+----+
|col1|col2|
+----+----+
|   A|   8|
|   B|   7|
+----+----+

Я хочу list из четных чисел от 0 до col2.

#Returns even numbers
def make_list(col):
    return list(map(int,[x for x in range(col+1) if x % 2 == 0]))
make_list = udf(make_list)

df = df.withColumn('list',make_list(col('col2')))
df.show()
+----+----+---------------+
|col1|col2|           list|
+----+----+---------------+
|   A|   8|[0, 2, 4, 6, 8]|
|   B|   7|   [0, 2, 4, 6]|
+----+----+---------------+
df.printSchema()
root
 |-- col1: string (nullable = true)
 |-- col2: long (nullable = true)
 |-- list: string (nullable = true)

Я получаюсписок, который я хочу, но список имеет тип string, а не int, как вы можете видеть в printschema выше.

Как я могу получить list типа int?Без int типа я не могу explode этот фрейм данных.

Любые идеи относительно того, как я могу получить list из integers?

Answer 1

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

Мы можем реализовать эту функцию, а затем использовать некоторые манипуляции со строками и регулярные выражения, чтобы получить желаемый результат, используя только функции API. Несмотря на то, что это более сложно, этот должен все же быть быстрее, чем udf.

import pyspark.sql.functions as f

def getEvenNumList(x):
    n = f.floor(x/2)
    return f.split(
        f.concat(
            f.lit("0,"), 
            f.regexp_replace(
                (2./81.*(-9*n+f.pow(10, (n+1))-10)).cast('int').cast('string'), 
                r"(?<=\d)(?=\d)", 
                ","
            )
        ),
        ","
    ).cast("array<int>")

df = df.withColumn("list", getEvenNumList(f.col("col2")))
df.show()
#+----+----+---------------+
#|col1|col2|           list|
#+----+----+---------------+
#|   A|   8|[0, 2, 4, 6, 8]|
#|   B|   7|   [0, 2, 4, 6]|
#+----+----+---------------+

df.printSchema()
#root
# |-- col1: string (nullable = true)
# |-- col2: long (nullable = true)
# |-- list: array (nullable = true)
# |    |-- element: integer (containsNull = true)

Объяснение

Количество элементов в вашем желаемом списке - один плюс пол col2, деленный на 2. (Плюс 1 для ведущего 0). Игнорируйте 0 на данный момент и позвольте n быть полом col2, деленным на 2.

Если вы объединили числа в вашем списке вместе (как вы можете использовать str.join), результирующее число будет дано выражением:

2*sum(i*10**(n-i) for i in range(1,n+1))

Используя Wolfram Alpha, вы можете вычислить уравнение замкнутой формы для этой суммы.

Получив это число, вы можете преобразовать его в строку с добавлением в начале 0.

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

Answer 2

Вам необходимо указать тип возврата udf; чтобы получить list из int, используйте ArrayType(IntegerType()):

from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType, IntegerType

# specify the return type as ArrayType(IntegerType())
make_list_udf = udf(make_list, ArrayType(IntegerType()))

df = df.withColumn('list',make_list_udf(col('col2')))
df.show()
+----+----+------------+                                                        
|col1|col2|        list|
+----+----+------------+
|   A|   6|[0, 2, 4, 6]|
|   B|   7|[0, 2, 4, 6]|
+----+----+------------+

df.printSchema()
root
 |-- col1: string (nullable = true)
 |-- col2: long (nullable = true)
 |-- list: array (nullable = true)
 |    |-- element: integer (containsNull = true)

Или, если вы используете spark 2.4, вы можете использовать новую функцию sequence:

values = [('A',8),('B',7)]
df = sqlContext.createDataFrame(values,['col1','col2'])

from pyspark.sql.functions import sequence, lit, col
df.withColumn('list', sequence(lit(0), col('col2'), step=lit(2))).show()
+----+----+---------------+
|col1|col2|           list|
+----+----+---------------+
|   A|   8|[0, 2, 4, 6, 8]|
|   B|   7|   [0, 2, 4, 6]|
+----+----+---------------+