PySpark Pass Index Column to element_at()

user_10864326 0
Viewed 1049

Can you pass a pyspark column with index values into the element_at functions?

I have a df similar to the following:

------------------
 foo   |   bar
------------------
  1    |   [3, 2, 1]
  2    |   [3, 3, 1]
  1    |   [8, 2, 1]

And I am trying to use the foo as the index to create the following:

 baz  |
--------
  3   |
  3   |
  8   |

I am trying to use the following with no success:

df.withColumn('baz', F.element_at('bar', F.col('foo'))
pyspark
2 Answers

This code will get you the desired result:

df.withColumn("baz", F.col("bar")[F.col("foo") - F.lit(1)]).show()

+---+---------+---+
|foo|      bar|baz|
+---+---------+---+
|  1|[3, 2, 1]|  3|
|  2|[3, 3, 1]|  3|
|  1|[8, 2, 1]|  8|
+---+---------+---+

I had to subtract one from foo because it looks like you're using 1-indexing. PySpark arrays are zero-indexed.

This is the more standard solution.

df.withColumn("baz", F.col("bar")[F.col("foo")]).show()

+---+---------+---+
|foo|      bar|baz|
+---+---------+---+
|  1|[3, 2, 1]|  2|
|  2|[3, 3, 1]|  1|
|  1|[8, 2, 1]|  2|
+---+---------+---+

Here's the code to make your sample data:

data = [(1, [3, 2, 1]), (2, [3, 3, 1]), (1, [8, 2, 1])]
df = spark.createDataFrame(data, ["foo", "bar"])
user_11251590

You could also use sql expression using .expr.

If foo column is already int( not bigint), there is no need to cast.

df.withColumn("baz", F.expr("""element_at(bar,cast(foo as int))""")).show()

#+---+---------+---+
#|foo|      bar|baz|
#+---+---------+---+
#|  1|[3, 2, 1]|  3|
#|  2|[3, 3, 1]|  3|
#|  1|[8, 2, 1]|  8|
#+---+---------+---+
user_101073890
Related