PySpark equivalent of function "typedLit" from Scala API

user_7220386 0
Viewed 3954

We have a function typedLit in Scala API for Spark to add the Array or Map as column value.

import org.apache.spark.sql.functions.typedLit
val df1 = Seq((1, 0), (2, 3)).toDF("a", "b")

df1.withColumn("seq", typedLit(Seq(1,2,3)))
    .show(truncate=false)

+---+---+---------+
|a  |b  |seq      |
+---+---+---------+
|1  |0  |[1, 2, 3]|
|2  |3  |[1, 2, 3]|
+---+---+---------+

I couldn't find the equivalent in PySpark. How can we create a column in PySpark with Array as a column value?

pyspark apache-spark-sql apache-spark scala arrays
3 Answers

There isn't an equivalent function in pyspark yet, but you can have an array column as shown below:

from pyspark.sql.functions import array, lit
df = sc.parallelize([[1,2], [3,4]]).toDF(['a', 'b'])
df.withColumn('seq', array([lit(i) for i in [1,2,3]])).show()

Output:

+---+---+---------+                                                             
|  a|  b|      seq|
+---+---+---------+
|  1|  2|[1, 2, 3]|
|  3|  4|[1, 2, 3]|
+---+---+---------+
user_21429940

Using expr and array looks the most elegant to me:

df = df.withColumn('seq', F.expr('array(1,2,3)'))

Test results:

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame([(1,0), (2,3)], ['a', 'b'])
df = df.withColumn('seq', F.expr('array(1,2,3)'))
df.show()
#  +---+---+---------+
#  |  a|  b|      seq|
#  +---+---+---------+
#  |  1|  0|[1, 2, 3]|
#  |  2|  3|[1, 2, 3]|
#  +---+---+---------+

Use F.expr('sequence(1,3)') if array numbers need to go in a sequence.

user_27535010

You can use .cast() directly after the lit() call to type the Column:

import pyspark.sql.functions as sf
from pyspark.sql.types import LongType

df1.withColumn("long", sf.lit(1).cast(LongType()))

The same works for array():

import pyspark.sql.functions as sf
from pyspark.sql.types import LongType, ArrayType
df1.withColumn("pirate", sf.array([sf.lit(x).cast(LongType()) for x in [1, 2, 3]]))
df1.withColumn("pirate", sf.array([sf.lit(x) for x in [1, 2, 3]]).cast(ArrayType(LongType())))

and if you really like text and typing but hate types, you could use:

df1.withColumn("pirate", sf.array(sf.lit("1"), sf.lit("2")).cast("array<int>"))

;)

PS Also consider using map with sf.lit instead of the for comprehension.

user_4597920
Related