How to change multiple column values to a constant with out specifying all column names?

Viewed 842

I have a very wide dataframe in spark .It has 80 columns so I want to set a column to 0 and rest to 1.

So the one I want to set to 1 I tried it with

df = df.withColumn("set_zero_column", lit(0))

and it worked.

Now I want to set the rest columns to 1. How do I do without specify all the 79 names ?

Any help is appreciated

3 Answers

Use select with a list comprehension:

from pyspark.sql.functions import lit

set_one_columns = [lit(1).alias(c) for c in df.columns if c != "set_zero_column"]
df = df.select(lit(0).alias("set_zero_column"), *set_one_columns)

If you needed to maintain the original column order, you could do:

cols = [lit(0).alias(c) if c == "set_zero_column" else lit(1).alias(c) for c in df.columns]
df = df.select(*cols)

I try to answer in Scala:

Example:

Method1:

//sample dataframe

val df=Seq(("a",1)).toDF("id","id1")

//filter req columns and add literal value

val cls=df.columns.map(x => if (x != "id1") (x,lit("1")) else (x,lit("0")))

//use foldLeft and add columns dynamically

val df2=cls.foldLeft(df){(df,cls) => df.withColumn(cls._1,cls._2)}

Result:

df2.show()
+---+---+
| id|id1|
+---+---+
|  1|  0|
+---+---+

Method2:Pault approach :)

val cls=df.columns.map( x => if (x !="id1") lit(1).alias(s"${x}") else lit(0).alias(s"${x}"))

Result:

df.select(cls:_*).show()
+---+---+
| id|id1|
+---+---+
|  1|  0|
+---+---+

I am still new to spark sql , although it may not be the most efficient way to handle this scenario but will be glad if it helps or can be further improved, This is how I was able to do it in java.

Step1: Create Sparksession and Load your file into daframe. code:

public void process() throws AnalysisException {
SparkSession session = new SparkSession.Builder()
.appName("Untyped Agregation on data frame")
.master("local")
.getOrCreate();

//Load the file that you need to compute. 

Dataset<Row> peopledf = session.read()
.option("header","true")
.option("delimiter"," ")
.csv("src/main/resources/Person.txt");

Output:

+--------+---+--------+
|    name|age|property|
+--------+---+--------+
|  Gaurav| 27|       1|
| Dheeraj| 30|       1|
|  Saloni| 26|       1|
|  Deepak| 30|       1|
|      Db| 25|       1|
|Praneeth| 24|       1|
|   jyoti| 26|       1|
+--------+---+--------+

Step2(optional):

Incase you require to provide constant value to any one column.

code:

//incase you require to chnage value for a single column.
Dataset<Row> peopledf1 = peopledf.withColumn("property",lit("0"));
peopledf1.show();

output:
+--------+---+--------+
|    name|age|property|
+--------+---+--------+
|  Gaurav| 27|       0|
| Dheeraj| 30|       0|
|  Saloni| 26|       0|
|  Deepak| 30|       0|
|      Db| 25|       0|
|Praneeth| 24|       0|
|   jyoti| 26|       0|
+--------+---+--------+

Step3:

Get the String array of all the column names in your data frame. code:

//Get the list of all the coloumns
String[] myStringArray = peopledf1.columns();

Step4: logic for filtering out the column among the array you dont want to provide constant value to and creating List of required columns names and lit("constsnt") for withColumns

code:

        //create two list one bieng names of columns you need to compute
        //other bieng same size(same number of element as that of column list)of 
        //lit("0") i.e constant
        //filter out the coloumn that you dont want to apply constant upon.
        List<String> myList = new ArrayList<String>();
        List<Column> myList1 = new ArrayList<Column>();
        for(String element : myStringArray){
            if(!(element.contains("name"))){
                myList.add(element);
                myList1.add(lit("0"));
            }
        }

Step5: Convert the List to Scala Seq as withColumns method requires that format of argument. code:

    //convert both list into scala Seq<Columns> and Seq<String> respectively.
    //Need to do this because withColumns method requires arguments in Seq form.
    //check scala doc for with columns
    Seq<Column> mySeq1 = convertListToSeq(myList1);
    Seq<String> mySeq= convertListToSeq1(myList);

code for convertListToSeq using JavaConverters:

   //Use JavaConverters to Convert List to Scala Seq using provided method below
    public Seq<String> convertListToSeq1(List<String> inputList) {
    return 
   JavaConverters.asScalaIteratorConverter(inputList.iterator()).asScala().
   toSeq();
   }

   public Seq<Column> convertListToSeq(List<Column> inputList) {
   return JavaConverters.asScalaIteratorConverter(inputList.iterator())
   .asScala().toSeq();
    }

Step6: Print Output to console code:

    //Display the required output on console.
    peopledf1.withColumns(mySeq,mySeq1).show();

output:

+--------+---+--------+
|    name|age|property|
+--------+---+--------+
|  Gaurav|  0|       0|
| Dheeraj|  0|       0|
|  Saloni|  0|       0|
|  Deepak|  0|       0|
|      Db|  0|       0|
|Praneeth|  0|       0|
|   jyoti|  0|       0|
+--------+---+--------+

Please do comment if code can be improved further.

Happy Learning, Gaurav

Related