CODEWITHSUNDEEP
pythonapache-sparkpysparkapache-spark-sql
xcen
xcen

Reputation: 692

Pyspark: adding a new column to dataframe based on the values in another dataframe using an udf

I have two pyspark dataframes and I am trying to add a new column to dataframe_2 (df_2) based on the values of dataframe_1. Dataframe_2 columns val_1 and val_2 should be row and column position of the dataframe_1.

dataframe_1

df_1 = sqlContext.createDataFrame([(0.78, 0.79, 0.45, 0.67, 0.88), (0.77, 0.79, 0.81, 0.82, 0.66), (0.99, 0.92, 0.94, 0.95, 0.91), (
    0.75, 0.53, 0.83, 0.73, 0.56), (0.77, 0.78, 0.99, 0.34, 0.67)], ["col_1", "col_2", "col_3", "col_4", "col_5"])

df_1.show()
+-----+-----+-----+-----+-----+
|col_1|col_2|col_3|col_4|col_5|
+-----+-----+-----+-----+-----+
| 0.78| 0.79| 0.45| 0.67| 0.88|
| 0.77| 0.79| 0.81| 0.82| 0.66|
| 0.99| 0.92| 0.94| 0.95| 0.91|
| 0.75| 0.53| 0.83| 0.73| 0.56|
| 0.77| 0.78| 0.99| 0.34| 0.67|
+-----+-----+-----+-----+-----+

dataframe_2

df_2 = sqlContext.createDataFrame([(34563, 435353424, 1, 2 ), (23524, 466344656, 2, 1), (52452, 263637236, 2, 5), (
   52334, 466633353, 2, 3), (66334, 563555578, 5, 4), (42552, 123445563, 5, 3), (72331, 413555213, 4, 3), (82311, 52355563, 2, 2)], ["id", "col_A", "val_1", "val_2"])
df_2.show()
+-----+---------+-----+-----+
|   id|    col_A|val_1|val_2|
+-----+---------+-----+-----+
|34563|435353424|    1|    2|
|23524|466344656|    2|    1|
|52452|263637236|    2|    5|
|52334|466633353|    2|    3|
|66334|563555578|    5|    4|
|42552|123445563|    5|    3|
|72331|413555213|    4|    3|
|82311| 52355563|    2|    2|
+-----+---------+-----+-----+

Goal: adding a new column to df_2 based on the values in df_1

I tried using creating an udf and got an error.

Expected output:

+-----+---------+-----+-----+---------------+
|   id|    col_A|val_1|val_2|value_from_df_1|
+-----+---------+-----+-----+---------------+
|34563|435353424|    1|    2|           0.79|
|23524|466344656|    2|    1|           0.77|
|52452|263637236|    2|    5|           0.66|
|52334|466633353|    2|    3|           0.94|
|66334|563555578|    5|    4|           0.34|
|42552|123445563|    5|    3|           0.99|
|72331|413555213|    4|    3|           0.83|
|82311| 52355563|    2|    2|           0.79|
+-----+---------+-----+-----+---------------+

My code:

from pyspark.sql import functions as F
import pyspark.sql.types as t

def add_data_to_table(table, value_1, value_2):
    return float(table.collect()[value_1-1][value_2-1])

select_data_from_table = F.udf(add_data_to_table, t.FloatType())
result_df = df_2.withColumn('value_from_df_1', select_data_from_table(df_1, df_2.val_1, df_2.val_2 ))
result_df.show()

Really appreciate it if someone can help. Thank you.

Upvotes: 2

Views: 279

Answers (1)

mck
mck

Reputation: 42342

Unlike pandas, Spark does not have concept of index, so you need to manually add an index. UDF is not appropriate here, because UDF operate on a row-by-row basis, not on the whole dataframe.

from pyspark.sql import functions as F, Window

df_1_id = df_1.withColumn(
    'row',
    F.row_number().over(Window.orderBy(F.monotonically_increasing_id()))
).select(
    'row',
    F.posexplode(F.array(*df_1.columns))
)

result = df_2.withColumn(
    'rowid',
    F.monotonically_increasing_id()
).join(
    df_1_id,
    (df_1_id.row == df_2.val_1) & (df_1_id.pos + 1 == df_2.val_2)
).orderBy('rowid').drop('rowid', 'row', 'pos')

result.show()
+-----+---------+-----+-----+----+
|   id|    col_A|val_1|val_2| col|
+-----+---------+-----+-----+----+
|34563|435353424|    1|    2|0.79|
|23524|466344656|    2|    1|0.77|
|52452|263637236|    2|    5|0.66|
|52334|466633353|    2|    3|0.81|
|66334|563555578|    5|    4|0.34|
|42552|123445563|    5|    3|0.99|
|72331|413555213|    4|    3|0.83|
|82311| 52355563|    2|    2|0.79|
+-----+---------+-----+-----+----+

Upvotes: 6

Related Questions