PySpark: Using column name in UDF and doing concatenation of column names based on a logic

user_11471365 0
Viewed 597

I have a Spark dataframe like this...

ID A B C D
id1 1 0 0 2
id2 0 3 0 1
id3 1 2 5 0
id4 4 0 0 1

I would like to have a new dataframe which is based on this logic...

  1. Accept any column which has positive values
  2. Concatenate their names

The result would then be like this...

ID NewColumn
id1 A,D
id2 B,D
id3 A,B,C
id4 A,D

My Effort:

A) For the first step I thought I will convert the integers into the name of the columns... So it will look like this...

ID A B C D
id1 A 0 0 D
id2 0 B 0 D
id3 A B C 0
id4 A 0 0 D

I was trying with a UDF but it did not work...

def CountSelect(colname, x):
  if x>0 :
    return colname
  else:
    return ""

countUDF = UserDefinedFunction(CountSelect, T.StringType())

cols = inoutDF.columns
cols.remove("ID")

intermediateDF = inputDF.select("ID", *(countDF(c, col(c)).alias(c) for c in cols))

BUT it is not working...

Can any of you help please?

B) Then I will use string concat function on all the columns

This part should be easier but still if you can combine both logic into a simpler working code, I would be really thankful to you.

Thanks a lot

pyspark apache-spark-sql apache-spark dataframe python
2 Answers

The idea is to mark rows across the columns which are positive and return the value of the respective column.

You can use reduce to mark the columns and create a new DataFrame and Finally use concat_ws to form the required value

A more concise solution as provided by @anky

Concise Solution -

sparkDF.withColumn("GreaterThanZero",F.concat_ws(",",*[F.when(F.col(col)>0,col) for col in to_concat]))\
.select("id","GreaterThanZero").show()

+---+---------------+
| id|GreaterThanZero|
+---+---------------+
|id1|            A,D|
|id2|            B,D|
|id3|          A,B,C|
|id4|            A,D|
+---+---------------+

Data Preparation

input_str = """
id1 1   0   0   2
id2 0   3   0   1
id3 1   2   5   0
id4 4   0   0   1
""".split()

input_values = list(map(lambda x: x.strip() if x.strip() != 'null' else None, input_str))

cols = list(map(lambda x: x.strip() if x.strip() != 'null' else None, "ID   A   B   C   D".split()))
            
n = len(input_values)
n_cols = 5

input_list = [tuple(input_values[i:i+n_cols]) for i in range(0,n,n_cols)]

sparkDF = sql.createDataFrame(input_list, cols)

sparkDF.show()

+---+---+---+---+---+
| ID|  A|  B|  C|  D|
+---+---+---+---+---+
|id1|  1|  0|  0|  2|
|id2|  0|  3|  0|  1|
|id3|  1|  2|  5|  0|
|id4|  4|  0|  0|  1|
+---+---+---+---+---+

Reduce

to_check = ['id','A','B','C','D']

sparkDF_marked = reduce(lambda df
                , x: df.withColumn(x,F.when(F.col(x) > 0 ,x).otherwise(None))\
                        if x != 'id' else df.withColumn(x,F.col(x)) \
                ,to_check, sparkDF
            )
                      
sparkDF_marked.show()

+---+----+----+----+----+
| id|   A|   B|   C|   D|
+---+----+----+----+----+
|id1|   A|null|null|   D|
|id2|null|   B|null|   D|
|id3|   A|   B|   C|null|
|id4|   A|null|null|   D|
+---+----+----+----+----+

Concat

to_concat = ['A','B','C','D']

sparkDF_marked.select(['id',F.concat_ws(',',*to_concat).alias('GreaterThanZero')]).show()

+---+---------------+
| id|GreaterThanZero|
+---+---------------+
|id1|            A,D|
|id2|            B,D|
|id3|          A,B,C|
|id4|            A,D|
+---+---------------+

The solution though works , but has some subtle nuances that you need to be careful off , especially the reduce code snippet and to_check and to_concat.

to_check can easily be replaced with - sparkDF.columns for the actual data , but do let me know the performance over the bigger dataset.

user_91089120

I have resolved it with UDF and let me post that here as well.

BTW: I modified a little bit and instead of making the final column as string with "," separated - I created a list and that serves my purpose better for the next steps in my project.

    from pyspark.sql.functions import lit, col, UserDefinedFunction, array
    import pyspark.sql.types as T

    def MakeOne(colname, x):
      return colname if x> 0 else None
    makeOneUDF = UserDefinedFunction(MakeOne, T.StringType())

    cols = inputDF.columns
    cols.remove("ID")

    def MakeList(arr):
      return [a for a in arr if a is not None]
    makeListUDF = UserDefinedFunction(MakeList, T.ArrayType(T.StringType()))

    outputDF = (inputDF.select("ID", *(makeOneUDF(lit(c), col(c)).alias(c) for c in cols)).withColumn("NewColumn", makeListUDF(array(*cols) )).select("ID", "NewColumn"))

Again, the NewColumn is of type Array Type or String Type and it stores the list of the columnnames.

| ID   | NewColumn|
|------|----------|
| id1  | [A,D]   |
| id2  | [B,D]   |
| id3  | [A,B,C] |
| id4  | [A,D]   |
user_114713650
Related