Unable to apply log function to a pyspark dataframe - numpy

Question

so I have a large dataset (about 1 TB+) where I have to do many operations for which I have thought of using pyspark for faster processing. Here are my imports:
import numpy as np
import pandas as pd
try:
import pyspark
from pyspark import SparkContext, SparkConf
from pyspark.sql import SparkSession, SQLContext
except ImportError as e:
raise ImportError('PySpark is not Configured')
print(f"PySpark Version : {pyspark.__version__}")
# Creating a Spark-Context
sc = SparkContext.getOrCreate(SparkConf().setMaster('local[*]'))
# Spark Builder
spark = SparkSession.builder \
.appName('MBLSRProcessor') \
.config('spark.executor.memory', '10GB') \
.getOrCreate()
# SQL Context - for SQL Query Executions
sqlContext = SQLContext(sc)
>> PySpark Version : 2.4.7
Now, I want to apply log10 function on two columns - For demonstrations, please consider this data:
data = spark.createDataFrame(pd.DataFrame({
"A" : [1, 2, 3, 4, 5],
"B" : [4, 3, 6, 1, 8]
}))
data.head(5)
>> [Row(A=1, B=4), Row(A=2, B=3), Row(A=3, B=6), Row(A=4, B=1), Row(A=5, B=8)]
This is what I require: log10(A + B) i.e. log10(6 + 4) = 1 for which I have made a function like this:
def add(a, b):
# this is for demonstration
return np.sum([a, b])
data = data.withColumn("ADD", add(data.A, data.B))
data.head(5)
>> [Row(A=1, B=4, ADD=5), Row(A=2, B=3, ADD=5), Row(A=3, B=6, ADD=9), Row(A=4, B=1, ADD=5), Row(A=5, B=8, ADD=13)]
But, I can't do the same for np.log10:
def np_log(a, b):
# actual function
return np.log10(np.sum([a, b]))
data = data.withColumn("LOG", np_log(data.A, data.B))
data.head(5)
TypeError Traceback (most recent call last)
<ipython-input-13-a5726b6c7dc2> in <module>
----> 1 data = data.withColumn("LOG", np_log(data.A, data.B))
2 data.head(5)
<ipython-input-12-0e020707faae> in np_log(a, b)
1 def np_log(a, b):
----> 2 return np.log10(np.sum([a, b]))
TypeError: loop of ufunc does not support argument 0 of type Column which has no callable log10 method

Answer 1

The best way to do this is to use native Spark functions:
import pyspark.sql.functions as F
import pandas as pd
data = spark.createDataFrame(pd.DataFrame({
"A" : [1, 2, 3, 4, 5],
"B" : [4, 3, 6, 1, 8]
}))
data = data.withColumn("LOG", F.log10(F.col('A') + F.col('B')))
But if you want, you can also use a UDF:
import pyspark.sql.functions as F
from pyspark.sql.types import FloatType
import numpy as np
import pandas as pd
data = spark.createDataFrame(pd.DataFrame({
"A" : [1, 2, 3, 4, 5],
"B" : [4, 3, 6, 1, 8]
}))
def udf_np_log(a, b):
# actual function
return float(np.log10(np.sum([a, b])))
np_log = F.udf(udf_np_log, FloatType())
data = data.withColumn("LOG", np_log(data.A, data.B))
+---+---+---------+
| A| B| LOG|
+---+---+---------+
| 1| 4| 0.69897|
| 2| 3| 0.69897|
| 3| 6|0.9542425|
| 4| 1| 0.69897|
| 5| 8|1.1139433|
+---+---+---------+
Interestingly it works for np.sum without UDF because I guess np.sum is just calling the + operator, which is valid for spark dataframe columns.