Pyspark collect_list and other aggregations with window functions on numeral columns with NaN values

Question

I want to convert a spark dataframe that looks like this below

sdf = spark.createDataFrame(pd.DataFrame([[1, 'name_value', None, 'square', 'polygon'],
                                          [2, 'name_value', 5.5, None, 'polygon']],
                                          columns=['id', 'name', 'area', 'sub_type', 'type']))

Fig 1

to look like Fig 2

essentially trying to de-duplicate some records be selecting the best available values (mostly just non-null) values for given name column

I am trying to use the window operations to get to it but am seeing some weird behavior with float column's NaN values

from pyspark.sql.functions import col
from pyspark.sql import functions as F, Window


name_window =  Window.partitionBy('name')

(sdf
 .withColumn('min_Id', F.min('id').over(name_window))
 .withColumn('prefered_area', F.first(col('area')).over(name_window.orderBy(col('area').desc_nulls_last())))
 .withColumn('prefered_sub_type', F.first('sub_type').over(name_window.orderBy(col('type').desc_nulls_last())))
 .withColumn('area_list', F.collect_list(col('area')).over(name_window.orderBy(col('area'))))
 .withColumn('sub_type_list', F.collect_list('sub_type').over(name_window.orderBy(col('type'))))
 .withColumn('prefered_type', F.first('type').over(name_window.orderBy(col('type').desc_nulls_last())))
 .toPandas())

Fig 3

See how the first and collect_list functions work fine for string columns but creates different results for float columns? like for the string column the prefered_sub_type column has the desired square value but the float column prefered_area is NaN. I also included the collect_list results to show the difference between the float and string columns.

Can someone please take a look and let me know what's the reason for this? and how can I get to my desired dataframe in Fig 2?

matkurek · Answer

I get different results copying and pasting your code, but:

Avoid ordering when using collect_list as it messes up results, just use window only:

F.collect_list(col('area')).over(name_window))

Use row_number to check for how exactly the sorting through window occurs to debug:

sdf.withColumn('prefered_area_rank', F.row_number().over(name_window.orderBy(col('area').desc_nulls_last())))

In order to align with your desired output I had to correct preferred column names you were using, also added the ranks to your sorts. Nulls don't show in lists, I'm using Spark 3.1.2 maybe that's why.

from pyspark.sql.functions import col
from pyspark.sql import functions as F, Window
import pandas as pd

sdf = spark.createDataFrame(pd.DataFrame([[1, 'name_value', None, 'square', 'polygon'],
                                          [2, 'name_value', 5.5, None, 'polygon']],
                                          columns=['id', 'name', 'area', 'sub_type', 'type']))

name_window =  Window.partitionBy('name')

(sdf
 .withColumn('min_Id', F.min('id').over(name_window))
 .withColumn('prefered_area', F.first(col('area')).over(name_window.orderBy(col('area').desc_nulls_last())))
 .withColumn('prefered_area_rank', F.row_number().over(name_window.orderBy(col('area').desc_nulls_last())))
 .withColumn('prefered_sub_type', F.first('sub_type').over(name_window.orderBy(col('sub_type').desc_nulls_last())))
 .withColumn('prefered_sub_type_rank', F.row_number().over(name_window.orderBy(col('sub_type').desc_nulls_last())))
 .withColumn('area_list', F.collect_list(col('area')).over(name_window))
 .withColumn('sub_type_list', F.collect_list('sub_type').over(name_window))
 .withColumn('prefered_type', F.first('type').over(name_window.orderBy(col('type').desc_nulls_last())))
 .toPandas())

Pyspark collect_list and other aggregations with window functions on numeral columns with NaN values

1 Answers