web-dev-qa-db-fra.com

Évitez l'impact sur les performances d'un mode de partition unique dans les fonctions de fenêtre Spark

Ma question est déclenchée par le cas d'utilisation du calcul des différences entre les lignes consécutives dans un spark dataframe.

Par exemple, j'ai:

>>> df.show()
+-----+----------+
|index|      col1|
+-----+----------+
|  0.0|0.58734024|
|  1.0|0.67304325|
|  2.0|0.85154736|
|  3.0| 0.5449719|
+-----+----------+

Si je choisis de les calculer à l'aide des fonctions "Fenêtre", je peux le faire comme ceci:

>>> winSpec = Window.partitionBy(df.index >= 0).orderBy(df.index.asc())
>>> import pyspark.sql.functions as f
>>> df.withColumn('diffs_col1', f.lag(df.col1, -1).over(winSpec) - df.col1).show()
+-----+----------+-----------+
|index|      col1| diffs_col1|
+-----+----------+-----------+
|  0.0|0.58734024|0.085703015|
|  1.0|0.67304325| 0.17850411|
|  2.0|0.85154736|-0.30657548|
|  3.0| 0.5449719|       null|
+-----+----------+-----------+

Question: J'ai partitionné explicitement la trame de données dans une seule partition. Quel en est l'impact sur les performances et, le cas échéant, pourquoi en est-il ainsi et comment pourrais-je l'éviter? Parce que lorsque je ne spécifie pas de partition, j'obtiens l'avertissement suivant:

16/12/24 13:52:27 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.
apache-sparkpysparkapache-spark-sqlpartitioningwindow-functions
15
Ytsen de Boer

Dans la pratique, l'impact sur les performances sera presque le même que si vous omettiez la clause partitionBy. Tous les enregistrements seront mélangés sur une seule partition, triés localement et itérés séquentiellement un par un.

La différence réside uniquement dans le nombre de partitions créées au total. Illustrons cela avec un exemple utilisant un ensemble de données simple avec 10 partitions et 1000 enregistrements:

df = spark.range(0, 1000, 1, 10).toDF("index").withColumn("col1", f.randn(42))

Si vous définissez un cadre sans partition par clause

w_unpart = Window.orderBy(f.col("index").asc())

et l'utiliser avec lag

df_lag_unpart = df.withColumn(
    "diffs_col1", f.lag("col1", 1).over(w_unpart) - f.col("col1")
)

il n'y aura qu'une seule partition au total:

df_lag_unpart.rdd.glom().map(len).collect()

[1000]

Par rapport à cette définition de trame avec index factice (simplifié un peu par rapport à votre code:

w_part = Window.partitionBy(f.lit(0)).orderBy(f.col("index").asc())

utilisera un nombre de partitions égal à spark.sql.shuffle.partitions:

spark.conf.set("spark.sql.shuffle.partitions", 11)

df_lag_part = df.withColumn(
    "diffs_col1", f.lag("col1", 1).over(w_part) - f.col("col1")
)

df_lag_part.rdd.glom().count()

11

avec une seule partition non vide:

df_lag_part.rdd.glom().filter(lambda x: x).count()

1

Malheureusement, il n'y a pas de solution universelle qui puisse être utilisée pour résoudre ce problème dans PySpark. C'est juste un mécanisme inhérent à l'implémentation combiné avec un modèle de traitement distribué.

Puisque la colonne index est séquentielle, vous pouvez générer une clé de partitionnement artificielle avec un nombre fixe d'enregistrements par bloc:

rec_per_block  = df.count() // int(spark.conf.get("spark.sql.shuffle.partitions"))

df_with_block = df.withColumn(
    "block", (f.col("index") / rec_per_block).cast("int")
)

et l'utiliser pour définir la spécification du cadre:

w_with_block = Window.partitionBy("block").orderBy("index")

df_lag_with_block = df_with_block.withColumn(
    "diffs_col1", f.lag("col1", 1).over(w_with_block) - f.col("col1")
)

Cela utilisera le nombre attendu de partitions:

df_lag_with_block.rdd.glom().count()

11

avec une distribution de données à peu près uniforme (nous ne pouvons pas éviter les collisions de hachage):

df_lag_with_block.rdd.glom().map(len).collect()

[0, 180, 0, 90, 90, 0, 90, 90, 100, 90, 270]

mais avec un certain nombre de lacunes sur les limites des blocs:

df_lag_with_block.where(f.col("diffs_col1").isNull()).count()

12

Les limites étant faciles à calculer:

from itertools import chain

boundary_idxs = sorted(chain.from_iterable(
    # Here we depend on sequential identifiers
    # This could be generalized to any monotonically increasing
    # id by taking min and max per block
    (idx - 1, idx) for idx in 
    df_lag_with_block.groupBy("block").min("index")
        .drop("block").rdd.flatMap(lambda x: x)
        .collect()))[2:]  # The first boundary doesn't carry useful inf.

vous pouvez toujours sélectionner:

missing = df_with_block.where(f.col("index").isin(boundary_idxs))

et remplissez-les séparément:

# We use window without partitions here. Since number of records
# will be small this won't be a performance issue
# but will generate "Moving all data to a single partition" warning
missing_with_lag = missing.withColumn(
    "diffs_col1", f.lag("col1", 1).over(w_unpart) - f.col("col1")
).select("index", f.col("diffs_col1").alias("diffs_fill"))

et join:

combined = (df_lag_with_block
    .join(missing_with_lag, ["index"], "leftouter")
    .withColumn("diffs_col1", f.coalesce("diffs_col1", "diffs_fill")))

pour obtenir le résultat souhaité:

mismatched = combined.join(df_lag_unpart, ["index"], "outer").where(
    combined["diffs_col1"] != df_lag_unpart["diffs_col1"]
)
assert mismatched.count() == 0
15
user6910411