Clés primaires avec Apache Spark

Échelle :

Si vous n'avez besoin que de numéros uniques, vous pouvez utiliser zipWithUniqueId et recréez DataFrame. Tout d'abord, quelques importations et données factices :

import sqlContext.implicits._
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.{StructType, StructField, LongType}

val df = sc.parallelize(Seq(
    ("a", -1.0), ("b", -2.0), ("c", -3.0))).toDF("foo", "bar")

Extraire le schéma pour une utilisation ultérieure :

val schema = df.schema

Ajouter un champ d'identifiant :

val rows = df.rdd.zipWithUniqueId.map{
   case (r: Row, id: Long) => Row.fromSeq(id +: r.toSeq)}

Créer un DataFrame :

val dfWithPK = sqlContext.createDataFrame(
  rows, StructType(StructField("id", LongType, false) +: schema.fields))

La même chose en Python :

from pyspark.sql import Row
from pyspark.sql.types import StructField, StructType, LongType

row = Row("foo", "bar")
row_with_index = Row(*["id"] + df.columns)

df = sc.parallelize([row("a", -1.0), row("b", -2.0), row("c", -3.0)]).toDF()

def make_row(columns):
    def _make_row(row, uid):
        row_dict = row.asDict()
        return row_with_index(*[uid] + [row_dict.get(c) for c in columns])
    return _make_row

f = make_row(df.columns)

df_with_pk = (df.rdd
    .zipWithUniqueId()
    .map(lambda x: f(*x))
    .toDF(StructType([StructField("id", LongType(), False)] + df.schema.fields)))

Si vous préférez un numéro consécutif, vous pouvez remplacer zipWithUniqueId avec zipWithIndex mais c'est un peu plus cher.

Directement avec DataFrame API :

(Universal Scala, Python, Java, R avec à peu près la même syntaxe)

Auparavant, j'avais manqué monotonicallyIncreasingId fonction qui devrait fonctionner correctement tant que vous n'avez pas besoin de numéros consécutifs :

import org.apache.spark.sql.functions.monotonicallyIncreasingId

df.withColumn("id", monotonicallyIncreasingId).show()
// +---+----+-----------+
// |foo| bar|         id|
// +---+----+-----------+
// |  a|-1.0|17179869184|
// |  b|-2.0|42949672960|
// |  c|-3.0|60129542144|
// +---+----+-----------+

Bien qu'utile monotonicallyIncreasingId est non déterministe. Non seulement les identifiants peuvent être différents d'une exécution à l'autre, mais sans astuces supplémentaires, ils ne peuvent pas être utilisés pour identifier les lignes lorsque les opérations suivantes contiennent des filtres.

Remarque :

Il est également possible d'utiliser rowNumber fonction fenêtre :

from pyspark.sql.window import Window
from pyspark.sql.functions import rowNumber

w = Window().orderBy()
df.withColumn("id", rowNumber().over(w)).show()

Malheureusement :

WARN Window :Aucune partition définie pour le fonctionnement de la fenêtre ! Déplacer toutes les données vers une seule partition peut entraîner une grave dégradation des performances.

Donc, à moins que vous n'ayez un moyen naturel de partitionner vos données et de garantir l'unicité, ce n'est pas particulièrement utile pour le moment.