find values closest to a list of values in pyspark

user_3211950 0
Viewed 285

Let's suppose to have this Pyspark dataframe:

x = np.random.randint(1, 100, 1000)
y = np.random.randint(1, 100, 1000)
z = np.random.randint(1, 100, 1000)

df = pd.DataFrame({'x': x, 'y': y, 'z': z})
spark_df = spark.createDataFrame(df)

Let's suppose to have this list of values:

lst = [10, 20, 30]

I would like to retrieve all the 3 (=len(lst)) rows of spark_df such that the difference between each value of lst and spark_df.x is the lowest. I would like to retrieve this three values as a spark dataframe. E.g:

+---+---+---+
|  x|  y|  z|
+---+---+---+
| 11| 32| 84|
| 22| 12| 38|
| 29| 14| 12|
+---+---+---+

In this case:

  • 11 is the closest spark_df.x value to 10
  • 22 is the closest spark_df.x value to 20
  • 29 is the closest spark_df.x value to 30

How can achieve this result in Pyspark 3+?

Note: this is only a toy example and the list of values could be in the order of thousands.

pyspark apache-spark dataframe python
2 Answers

Step 1: add columns with the difference of the elements of lst and the x values to the dataframe:

from pyspark.sql import functions as F

diffs = [F.abs(F.col("x") - F.lit(c)).alias(f"diff_{c}") for c in lst]
df_with_diffs = spark_df.select("*", *diffs)

+---+---+---+-------+-------+-------+
|  x|  y|  z|diff_10|diff_20|diff_30|
+---+---+---+-------+-------+-------+
| 15| 34| 20|      5|      5|     15|
| 12| 45| 24|      2|      8|     18|
| 86| 49| 13|     76|     66|     56|
+---+---+---+-------+-------+-------+
only showing top 3 rows

Step 2: Collect the minimal values for each of the diff columns and select the respective rows:

mins=df_with_diffs.select(*[F.min(f"diff_{c}") for c in lst]).first()

filter=" or ".join([f"(diff_{c} = {mins[i]})" for i,c in enumerate(lst)])
df_with_diffs.filter(filter).select(spark_df.columns).show()

+---+---+---+
|  x|  y|  z|
+---+---+---+
| 12| 45| 24|
| 22| 28| 58|
| 27| 96| 36|
+---+---+---+

Step 2 (original answer): use min_by for each of the newly created columns to find the row with the minimal difference. For each value of lst this returns one row. All these rows are then unioned.

agg_cols = [[F.expr(f"min_by({c}, diff_{val})").alias(c) for c in spark_df.columns] 
  for val in lst]

import functools
result = functools.reduce(lambda a,b: a.union(df_with_diffs.agg(*b)), agg_cols[1:], 
          df_with_diffs.agg(*agg_cols[0]))
result.show()
user_21298010

After some experiments on my own, I propose this solution. At first, I wanted avoid pandas_udf function, but, none the less, it seems elegant, pythonic and effective.

Step 1: Create a pandas_udf in order to obtain a list of minimum values:

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType, ArrayType
import pyspark.sql.functions as func

bdc_values = sc.broadcast([10, 20, 30])

@pandas_udf(ArrayType(IntegerType()))
def get_x_min(x: pd.Series) -> ArrayType(IntegerType()):
   values = bdc_values.value
   return [x.iloc[(x - v).abs().argmin()] for v in values]

Step 2: Apply the pandas_udf to dataframe and collect the values, in this way:

mins = spark_df.agg(get_x_min('x')).first()[0]

Step 3: Filter the spark_df with isin function and then eventually deduplicate:

result = spark_df.filter(func.col('x').isin(mins)).dropDuplicates(['x'])
user_32119500
Related