web-dev-qa-db-fra.com

comment ajouter un identifiant de ligne dans des cadres de données pySpark

J'ai un fichier csv; que je convertis en DataFrame (df) dans pyspark; après une certaine transformation; Je veux ajouter une colonne en df; qui doit être un simple identifiant de ligne (en commençant de 0 ou 1 à N).

J'ai converti df en rdd et j'utilise "zipwithindex". J'ai converti le rdd résultant en df. cette approche fonctionne, mais elle a généré 250 000 tâches et prend beaucoup de temps à exécuter. Je me demandais s'il y avait une autre façon de le faire qui prend moins de temps d'exécution.

voici un extrait de mon code; le fichier csv que je traite est GRAND; contient des milliards de lignes.

debug_csv_rdd = (sc.textFile("debug.csv")
  .filter(lambda x: x.find('header') == -1)
  .map(lambda x : x.replace("NULL","0")).map(lambda p: p.split(','))
  .map(lambda x:Row(c1=int(x[0]),c2=int(x[1]),c3=int(x[2]),c4=int(x[3]))))

debug_csv_df = sqlContext.createDataFrame(debug_csv_rdd)
debug_csv_df.registerTempTable("debug_csv_table")
sqlContext.cacheTable("debug_csv_table")

r0 = sqlContext.sql("SELECT c2 FROM debug_csv_table WHERE c1 = 'str'")
r0.registerTempTable("r0_table")

r0_1 = (r0.flatMap(lambda x:x)
    .zipWithIndex()
    .map(lambda x: Row(c1=x[0],id=int(x[1]))))

r0_df=sqlContext.createDataFrame(r0_2)
r0_df.show(10) 
17
ankit patel

Vous pouvez également utiliser une fonction du package sql. Il générera un identifiant unique, mais il ne sera pas séquentiel car cela dépend du nombre de partitions. Je pense qu'il est disponible en Spark 1.5 +

from pyspark.sql.functions import monotonicallyIncreasingId

# This will return a new DF with all the columns + id
res = df.withColumn("id", monotonicallyIncreasingId())

Modifier: 19/1/2017

Comme commenté par @ Sean

Utilisez monotonically_increasing_id() à la place de Spark 1.6 et plus

54
Arkadi T