PySpark-numpy interoperability

Question

I am working on a dataset where most of the columns are 1D sequences of doubles, like the following one:

source_data = [
    Row(ID="F0", P1=[-1.0, -2.0, -3.0], P2=[-4.0, -5.0, -6.0]),
    Row(ID="F1", P1=[1.0, 2.0, 3.0], P2=[4.0, 5.0, 6.0]),
]
df = spark.createDataFrame(source_data)

which looks like:

+---+------------------+------------------+
| ID|                P1|                P2|
+---+------------------+------------------+
| F0|[-1.0, -2.0, -3.0]|[-4.0, -5.0, -6.0]|
| F1|   [1.0, 2.0, 3.0]|   [4.0, 5.0, 6.0]|
+---+------------------+------------------+

The columns P1 and P2 have spark type ArrayType(DoubleType()). In my real dataset, I have hundreds of columns of different length.

My goal is to execute mathematical operations such as mean, median, quantiles, fft, etc on these sequences in a distributed manner on a cluster. My approach is to wrap numpy functions as follows. First, I "hide" the numpy types which do not seem to be supported by Spark:

# wrap functions to "hide" numpy data types
def py_median(param_val):
    param_val = np.array(param_val)
    param_median = np.median(param_val)
    return float(param_median)

def py_abs(param_val):
    param_val = np.array(param_val)
    param_abs = np.abs(param_val)
    return param_abs.tolist()

Then, I transform the functions to pySpark UDFs:

# wrap functions to operate as udfs
median = F.udf(py_median, DoubleType())
abs = F.udf(py_abs, ArrayType(DoubleType()))

Finally, I can use them with pySpark to do the work:

df_processed = (
    df.withColumn("P1_mean", median(col("P1")))
    .withColumn("P2_abs", abs(col("P2")))
)

Note that I have to convert the result type to simple python types such as:

python float (functions returning a scalar)
python list of float (functions returning a 1D numpy array) otherwise I get Spark errors

I need different wrapping strategies to handle functions returning a complex array (e.g. the fft) or multiple return value. Is there a cleaner way to get my job done? Perhaps using Pandas UDFs?

Thanks in advance,

Marco

PySpark-numpy interoperability

0 Answers