web-dev-qa-db-fra.com

Modifier la propriété nullable de la colonne dans spark dataframe

Je crée manuellement une trame de données pour certains tests. Le code pour le créer est:

case class input(id:Long, var1:Int, var2:Int, var3:Double)
val inputDF = sqlCtx
  .createDataFrame(List(input(1110,0,1001,-10.00),
    input(1111,1,1001,10.00),
    input(1111,0,1002,10.00)))

Le schéma ressemble donc à ceci:

root
 |-- id: long (nullable = false)
 |-- var1: integer (nullable = false)
 |-- var2: integer (nullable = false)
 |-- var3: double (nullable = false)

Je veux faire 'nullable = true' pour chacune de ces variables. Comment puis-je le déclarer dès le début ou le basculer dans une nouvelle trame de données après sa création?

scalaapache-sparkspark-dataframe
21
J Calbreath

Répondre

Avec les importations

import org.Apache.spark.sql.types.{StructField, StructType}
import org.Apache.spark.sql.{DataFrame, SQLContext}
import org.Apache.spark.{SparkConf, SparkContext}

vous pouvez utiliser

/**
 * Set nullable property of column.
 * @param df source DataFrame
 * @param cn is the column name to change
 * @param nullable is the flag to set, such that the column is  either nullable or not
 */
def setNullableStateOfColumn( df: DataFrame, cn: String, nullable: Boolean) : DataFrame = {

  // get schema
  val schema = df.schema
  // modify [[StructField] with name `cn`
  val newSchema = StructType(schema.map {
    case StructField( c, t, _, m) if c.equals(cn) => StructField( c, t, nullable = nullable, m)
    case y: StructField => y
  })
  // apply new schema
  df.sqlContext.createDataFrame( df.rdd, newSchema )
}

directement.

Vous pouvez également rendre la méthode disponible via le modèle de bibliothèque "pimp my library" (voir mon SO post Quelle est la meilleure façon de définir des méthodes personnalisées sur un DataFrame? ), de sorte que vous puissiez appeler

val df = ....
val df2 = df.setNullableStateOfColumn( "id", true )

Éditer

Solution alternative 1

Utilisez une version légèrement modifiée de setNullableStateOfColumn

def setNullableStateForAllColumns( df: DataFrame, nullable: Boolean) : DataFrame = {
  // get schema
  val schema = df.schema
  // modify [[StructField] with name `cn`
  val newSchema = StructType(schema.map {
    case StructField( c, t, _, m) ⇒ StructField( c, t, nullable = nullable, m)
  })
  // apply new schema
  df.sqlContext.createDataFrame( df.rdd, newSchema )
}

Solution alternative 2

Définissez explicitement le schéma. (Utilisez la réflexion pour créer une solution plus générale)

configuredUnitTest("Stackoverflow.") { sparkContext =>

  case class Input(id:Long, var1:Int, var2:Int, var3:Double)

  val sqlContext = new SQLContext(sparkContext)
  import sqlContext.implicits._


  // use this to set the schema explicitly or
  // use refelection on the case class member to construct the schema
  val schema = StructType( Seq (
    StructField( "id", LongType, true),
    StructField( "var1", IntegerType, true),
    StructField( "var2", IntegerType, true),
    StructField( "var3", DoubleType, true)
  ))

  val is: List[Input] = List(
    Input(1110, 0, 1001,-10.00),
    Input(1111, 1, 1001, 10.00),
    Input(1111, 0, 1002, 10.00)
  )

  val rdd: RDD[Input] =  sparkContext.parallelize( is )
  val rowRDD: RDD[Row] = rdd.map( (i: Input) ⇒ Row(i.id, i.var1, i.var2, i.var3))
  val inputDF = sqlContext.createDataFrame( rowRDD, schema ) 

  inputDF.printSchema
  inputDF.show()
}
38
Martin Senne

C'est une réponse tardive, mais je voulais donner une solution alternative pour les gens qui viennent ici. Vous pouvez automatiquement rendre un DataFrameColumn nullable dès le début par la modification suivante de votre code:

case class input(id:Option[Long], var1:Option[Int], var2:Int, var3:Double)
val inputDF = sqlContext
  .createDataFrame(List(input(Some(1110),Some(0),1001,-10.00),
    input(Some(1111),Some(1),1001,10.00),
    input(Some(1111),Some(0),1002,10.00)))
inputDF.printSchema

Cela donnera:

root
 |-- id: long (nullable = true)
 |-- var1: integer (nullable = true)
 |-- var2: integer (nullable = false)
 |-- var3: double (nullable = false)

defined class input
inputDF: org.Apache.spark.sql.DataFrame = [id: bigint, var1: int, var2: int, var3: double]

Essentiellement, si vous déclarez un champ en tant que Option en utilisant Some([element]) ou None comme entrées réelles, ce champ peut être annulé. Sinon, le champ ne sera pas annulable. J'espère que ça aide!

14
Sidd Singal

Version plus compacte de la définition du paramètre nullable de toutes les colonnes

Au lieu de case StructField( c, t, _, m) ⇒ StructField( c, t, nullable = nullable, m) on peut utiliser _.copy(nullable = nullable). Ensuite, toute la fonction peut s'écrire:

def setNullableStateForAllColumns( df: DataFrame, nullable: Boolean) : DataFrame = {
  df.sqlContext.createDataFrame(df.rdd, StructType(df.schema.map(_.copy(nullable = nullable))))
}
7
matemaciek

Une autre option, si vous devez modifier le cadre de données sur place et que la recréation est impossible, vous pouvez faire quelque chose comme ceci:

.withColumn("col_name", when(col("col_name").isNotNull, col("col_name")).otherwise(lit(null)))

Spark pensera alors que cette colonne peut contenir null, et la nullité sera définie sur true. Vous pouvez également utiliser udf, pour encapsuler vos valeurs dans Option. Fonctionne bien même pour les cas de streaming.

5
Rayan Ral

Utilisez simplement Java.lang.Integer au lieu de scala.Int dans votre classe de cas.

case class input(id:Long, var1:Java.lang.Integer , var2:Java.lang.Integer , var3:Java.lang.Double)
2
echo

Merci Martin Senne . Juste un petit ajout. Dans le cas de types de structures internes, vous devrez peut-être définir nullable récursivement, comme ceci:

def setNullableStateForAllColumns(df: DataFrame, nullable: Boolean): DataFrame = {
    def set(st: StructType): StructType = {
      StructType(st.map {
        case StructField(name, dataType, _, metadata) =>
          val newDataType = dataType match {
            case t: StructType => set(t)
            case _ => dataType
          }
          StructField(name, newDataType, nullable = nullable, metadata)
      })
    }

    df.sqlContext.createDataFrame(df.rdd, set(df.schema))
  }
0
skotlov