PySpark how to use col function that refer's to another column's value

user_10033921 0
Viewed 739

I have a dataframe that has below columns

field1 , field2 , field3 , field_name

sample data

"a1", "b1", "c1", "field1"
"a2", "b2", "c2", "field2"
"a3", "b3", "c3", "field3"

I want to add new column"fieldvalue" to the data frame, so that it contains value in the column that corresponds to the content of the column "fieldname"

so the first row above will have fieldvalue = "a1", since fieldname contains "field1"

the output data frame should look like

field1, field2, field3 , fieldname, fieldvalue

data

"a1", "b1", "c1", "field1", "a1"
"a2", "b2", "c2", "field2", "b2"
"a3", "b3", "c3", "field3", "c3"

I tried to use below syntax

df1 = df1.withColumn("fieldValue", func.col(func.col("fieldName")))

But it fails with below error, since func.col expects a column , and not a string

Method col([class org.apache.spark.sql.Column]) does not exist

pyspark python
1 Answers

Using col you can only use a constant value as parameter that does not depend on the values of the actual row. The reason for this that the DAG for the execution of this function is created before Spark even looks at the data. col simply picks a column from the dataframe and this column has to be the same for all rows.

Instead each row can be converted into an array, then the index within the array of the required column is calculated using find_in_set and finally the required value is taken from the array:

from pyspark.sql import functions as F

df.withColumn("array", F.array(df.columns)) \
  .withColumn("index", F.expr(f"find_in_set(field_name, '{','.join(df.columns)}') - 1")) \
  .withColumn("fieldvalue", F.expr("array[index]")) \
  .show()

Output:

+------+------+------+----------+--------------------+-----+----------+
|field1|field2|field3|field_name|               array|index|fieldvalue|
+------+------+------+----------+--------------------+-----+----------+
|    a1|    b1|    c1|    field1|[a1, b1, c1, field1]|    0|        a1|
|    a2|    b2|    c2|    field2|[a2, b2, c2, field2]|    1|        b2|
|    a3|    b3|    c3|    field3|[a3, b3, c3, field3]|    2|        c3|
+------+------+------+----------+--------------------+-----+----------+

The same result could be achieved using an udf, but Spark SQL functions are usually faster than an udf, especially when using Python.

Edit: if array does not work, the same idea can be used with map:

m = "map(" + ",".join([f"'{c}', {c}" for c in df.columns]) + ")"

df.withColumn("map", F.expr(m)) \
  .withColumn("fieldvalue", F.expr("map[field_name]")) \
  .show()
user_21298010
Related