web-dev-qa-db-fra.com

Spark Sélection de trame de données basée sur l'index de colonne

Comment sélectionner toutes les colonnes d'une trame de données qui a certains index dans Scala?

Par exemple, si une trame de données comporte 100 colonnes et que je souhaite extraire uniquement des colonnes (10,12,13,14,15), comment faire de même?

Ci-dessous sélectionne toutes les colonnes du dataframe df qui a le nom de colonne mentionné dans les noms de colonnes du tableau:

df = df.select(colNames.head,colNames.tail: _*)

S'il y a un tableau colNos similaire qui a

colNos = Array(10,20,25,45)

Comment puis-je transformer ce qui précède df.select pour récupérer uniquement ces colonnes aux index spécifiques.

8
Vikas J

Vous pouvez map sur columns:

import org.Apache.spark.sql.functions.col

df.select(colNos map df.columns map col: _*)

ou:

df.select(colNos map (df.columns andThen col): _*)

ou:

df.select(colNos map (col _ compose df.columns): _*)

Toutes les méthodes présentées ci-dessus sont équivalentes et n'imposent pas de pénalité de performance. Cartographie suivante:

colNos map df.columns 

est juste un accès local Array ( accès à temps constant pour chaque index ) et en choisissant entre String ou Column une variante basée sur select n'affecte pas le plan d'exécution:

val df = Seq((1, 2, 3 ,4, 5, 6)).toDF

val colNos = Seq(0, 3, 5)

df.select(colNos map df.columns map col: _*).explain
== Physical Plan ==
LocalTableScan [_1#46, _4#49, _6#51]
df.select("_1", "_4", "_6").explain
== Physical Plan ==
LocalTableScan [_1#46, _4#49, _6#51]
9
user6910411

La réponse de @ user6910411 ci-dessus fonctionne comme un charme et le nombre de tâches/plan logique est similaire à mon approche ci-dessous. MAIS mon approche est un peu plus rapide.
Donc,
Je vous suggère d'aller avec le column names plutôt que column numbers. Column names sont beaucoup plus sûrs et beaucoup plus légers que d'utiliser numbers. Vous pouvez utiliser la solution suivante:

val colNames = Seq("col1", "col2" ...... "col99", "col100")

val selectColNames = Seq("col1", "col3", .... selected column names ... )

val selectCols = selectColNames.map(name => df.col(name))

df = df.select(selectCols:_*)

Si vous hésitez à écrire les 100 noms de colonne, il existe également une méthode de raccourci

val colNames = df.schema.fieldNames
3
Ramesh Maharjan