web-dev-qa-db-fra.com

Filtrage d'un DataFrame Pyspark avec une clause IN de type SQL

Je veux filtrer un DataFrame Pyspark avec une clause IN de type SQL, comme dans

sc = SparkContext()
sqlc = SQLContext(sc)
df = sqlc.sql('SELECT * from my_df WHERE field1 IN a')

a est le tuple (1, 2, 3). Je reçois cette erreur:

Java.lang.RuntimeException: [1.67] échec: `(('' attendu mais identificateur a trouvé

ce qui revient à dire qu'il attendait quelque chose comme '(1, 2, 3)' au lieu de . Le problème est que je ne peux pas écrire manuellement les valeurs dans un fichier car elles sont extraites d'un autre travail.

Comment pourrais-je filtrer dans ce cas?

25
mar tin

La chaîne que vous transmettez à SQLContext a été évaluée dans le cadre de l'environnement SQL. Cela ne rend pas compte de la fermeture. Si vous voulez passer une variable, vous devrez le faire explicitement en utilisant le formatage de chaîne:

df = sc.parallelize([(1, "foo"), (2, "x"), (3, "bar")]).toDF(("k", "v"))
df.registerTempTable("df")
sqlContext.sql("SELECT * FROM df WHERE v IN {0}".format(("foo", "bar"))).count()
##  2 

Évidemment, ce n'est pas quelque chose que vous utiliseriez dans un "vrai" environnement SQL pour des raisons de sécurité, mais cela ne devrait pas avoir d'importance ici.

En pratique, DataFrame DSL est un choix judicieux lorsque vous souhaitez créer des requêtes dynamiques:

from pyspark.sql.functions import col

df.where(col("v").isin({"foo", "bar"})).count()
## 2

Il est facile de construire et de composer et gère tous les détails de HiveQL/Spark SQL pour vous.

46
zero323

réitérant ce que @ zero323 a mentionné ci-dessus: nous pouvons également faire la même chose en utilisant une liste (pas seulement set) comme ci-dessous

from pyspark.sql.functions import col

df.where(col("v").isin(["foo", "bar"])).count()
16
braj

Juste un petit ajout/mise à jour:

choice_list = ["foo", "bar", "jack", "joan"]

Si vous souhaitez filtrer votre cadre de données "df", de manière à conserver les lignes en fonction de la colonne "v" en prenant uniquement les valeurs de choice_list,

df_filtered = df.where( ( col("v").isin (choice_list) ) )
1
shwetabharti

Une approche légèrement différente qui a fonctionné pour moi consiste à filtrer avec une fonction de filtre personnalisée.

def filter_func(a):
"""wrapper function to pass a in udf"""
    def filter_func_(col):
    """filtering function"""
        if col in a.value:
            return True

    return False

return udf(filter_func_, BooleanType())

# Broadcasting allows to pass large variables efficiently
a = sc.broadcast((1, 2, 3))
df = my_df.filter(filter_func(a)(col('field1'))) \
0
Alex_Gidiotis