Aplicando condição quando somente quando a coluna existe no dataframe

Question

Estou usando spark-sql-2.4.1v com java8. Tenho um cenário em que preciso realizar determinada operação se as colunas estiverem presentes na lista de colunas do dataframe fornecida

Eu tenho o quadro de dados de amostra conforme abaixo, as colunas do quadro de dados seriam diferentes com base na consulta externa executada na tabela de banco de dados.

val data = List(
  ("20", "score", "school", "2018-03-31", 14 , 12 , 20),
  ("21", "score", "school", "2018-03-31", 13 , 13 , 21),
  ("22", "rate", "school", "2018-03-31", 11 , 14, 22),
  ("21", "rate", "school", "2018-03-31", 13 , 12, 23)
 )

val df = data.toDF("id", "code", "entity", "date", "column1", "column2" ,"column3"..."columnN")

conforme mostrado acima, as colunas de "dados" do dataframe não são fixas e variam e teriam "coluna1", "coluna2", "coluna3" ... "colunaN" ...

Então, dependendo da disponibilidade da coluna, eu preciso realizar algumas operações para a mesma, estou tentando usar a "cláusula quando", quando uma coluna está presente, então eu tenho que executar determinada operação na coluna especificada, caso contrário, passar para a próxima operação.

Estou tentando a seguir duas maneiras de usar "when-cluase"

Primeira via:

 Dataset<Row> resultDs =  df.withColumn("column1_avg", 
                     when( df.schema().fieldNames().contains(col("column1")) , avg(col("column1"))))
                     )
 

Segunda via:

  Dataset<Row> resultDs =  df.withColumn("column2_sum", 
                     when( df.columns().contains(col("column2")) , sum(col("column1"))))
                     )

Erro:

Não é possível invocar contém (coluna) no tipo de matriz String []

então, como lidar com esse cenário usando o código java8?

Answer 1

Você pode criar uma coluna com todos os nomes das colunas. então você pode verificar se a coluna está presente ou não e processar se ela está disponível-

 df.withColumn("columns_available", array(df.columns.map(lit): _*))
      .withColumn("column1_org",
      when( array_contains(col("columns_available"),"column1") , col("column1")))
      .withColumn("x",
        when( array_contains(col("columns_available"),"column4") , col("column1")))
      .withColumn("column2_new",
        when( array_contains(col("columns_available"),"column2") , sqrt("column2")))
      .show(false)