web-dev-qa-db-fra.com

Fonction Pyspark Dataframe Apply sur deux colonnes

Supposons que j'ai deux PySpark DataFrames df1 Et df2.

df1=   'a' 
        1    
        2    
        5    

df2=   'b'
        3
        6

Et je veux trouver la valeur df2['b'] La plus proche pour chaque df1['a'], Et ajouter les valeurs les plus proches en tant que nouvelle colonne dans df1.

En d'autres termes, pour chaque valeur x dans df1['a'], Je veux trouver un y qui réalise min(abx(x-y)) pour tous y in df2['b'] ( note: peut supposer qu'il n'y a qu'un seul y qui peut atteindre la distance minimale), et le résultat serait

'a'    'b'
 1      3
 2      3
 5      6

J'ai essayé le code suivant pour créer d'abord une matrice de distance (avant de trouver les valeurs atteignant la distance minimale):

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf

def dict(x,y):
    return abs(x-y)
udf_dict = udf(dict, IntegerType())

sql_sc = SQLContext(sc)
udf_dict(df1.a, df2.b)

qui donne

Column<PythonUDF#dist(a,b)>

J'ai essayé

sql_sc.CreateDataFrame(udf_dict(df1.a, df2.b))

qui s'exécute pour toujours sans donner d'erreur/sortie.

Mes questions sont:

  1. Comme je suis nouveau sur Spark, ma façon de construire le DataFrame de sortie est-elle efficace? (Ma façon serait de créer une matrice de distance pour toutes les valeurs a et b d'abord, puis de trouver celle min)
  2. Quel est le problème avec la dernière ligne de mon code et comment y remédier?
pysparkspark-dataframepyspark-sql
8
Chianti5

En commençant par votre deuxième question - vous pouvez appliquer udf uniquement à la trame de données existante, je pense que vous pensiez à quelque chose comme ceci:

>>> df1.join(df2).withColumn('distance', udf_dict(df1.a, df2.b)).show()
+---+---+--------+
|  a|  b|distance|
+---+---+--------+
|  1|  3|       2|
|  1|  6|       5|
|  2|  3|       1|
|  2|  6|       4|
|  5|  3|       2|
|  5|  6|       1|
+---+---+--------+

Mais il existe un moyen plus efficace d'appliquer cette distance, en utilisant abs interne:

>>> from pyspark.sql.functions import abs
>>> df1.join(df2).withColumn('distance', abs(df1.a -df2.b))

Ensuite, vous pouvez trouver des nombres correspondants en calculant:

>>> distances = df1.join(df2).withColumn('distance', abs(df1.a -df2.b))
>>> min_distances = distances.groupBy('a').agg(min('distance').alias('distance'))
>>> distances.join(min_distances, ['a', 'distance']).select('a', 'b').show()
+---+---+                                                                       
|  a|  b|
+---+---+
|  5|  6|
|  1|  3|
|  2|  3|
+---+---+
8
Mariusz