web-dev-qa-db-fra.com

Clés primaires avec Apache Spark

J'ai une connexion JDBC avec Apache Spark et PostgreSQL et je veux insérer des données dans ma base de données. Lorsque j'utilise le mode append, je dois spécifier id pour chaque DataFrame.Row. Existe-t-il un moyen pour Spark de créer des clés primaires?

databasepostgresqlhadoopapache-spark
23
Nhor

Scala:

Si vous n'avez besoin que de numéros uniques, vous pouvez utiliser zipWithUniqueId et recréer DataFrame. 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'identification:

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

Créez DataFrame:

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

La même chose dans 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 par zipWithIndex mais c'est un peu plus cher.

directement avec l'API DataFrame:

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

Auparavant, j'ai manqué la fonction monotonicallyIncreasingId qui devrait fonctionner très bien 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 n'est pas déterministe. Non seulement les identifiants peuvent être différents d'une exécution à l'autre, mais sans astuces supplémentaires 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 la fonction de fenêtre rowNumber:

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

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

Malheureusement:

Fenêtre WARN: aucune partition définie pour le fonctionnement de la fenêtre! Le déplacement de toutes les données vers une seule partition peut entraîner une grave dégradation des performances.

Donc, sauf si vous avez un moyen naturel de partitionner vos données et de vous assurer que l'unicité n'est pas particulièrement utile pour le moment.

37
zero323
from pyspark.sql.functions import monotonically_increasing_id

df.withColumn("id", monotonically_increasing_id()).show()

Notez que le 2ème argument de df.withColumn est monotonically_increasing_id () et non monotonically_increasing_id.

9
Allyn

J'ai trouvé que la solution suivante était relativement simple pour le cas où zipWithIndex () est le comportement souhaité, c'est-à-dire pour ceux qui souhaitent des entiers consécutifs.

Dans ce cas, nous utilisons pyspark et comptons sur la compréhension du dictionnaire pour mapper l'objet de ligne d'origine à un nouveau dictionnaire qui correspond à un nouveau schéma comprenant l'index unique.

# read the initial dataframe without index
dfNoIndex = sqlContext.read.parquet(dataframePath)
# Need to Zip together with a unique integer

# First create a new schema with uuid field appended
newSchema = StructType([StructField("uuid", IntegerType(), False)]
                       + dfNoIndex.schema.fields)
# Zip with the index, map it to a dictionary which includes new field
df = dfNoIndex.rdd.zipWithIndex()\
                      .map(lambda (row, id): {k:v
                                              for k, v
                                              in row.asDict().items() + [("uuid", id)]})\
                      .toDF(newSchema)
3
rocconnick