Spark dataframe to numpy array via udf or without collecting to driver

Question

Real life df is a massive dataframe that cannot be loaded into driver memory. Can this be done using regular or pandas udf?

# Code to generate a sample dataframe

from pyspark.sql import functions as F
from pyspark.sql.types import *
import pandas as pd
import numpy as np

sample = [['123',[[0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1], [0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
      ['345',[[1,0,0,0,0,1,1,1,0,1,1,0,1,0,0,0,1,1,1,1,1,1], [0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
      ['425',[[1,1,0,0,0,1,0,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1],[0,1,0,0,0,1,1,1,1,1,1,0,1,0,0,0,1,1,1,1,1,1]]],
      ]

df = spark.createDataFrame(sample,["id", "data"])

Here's the logic that needs to be parallelized without relying on driver memory.

Input: Spark dataframe Output: numpy array to be fed into horovod (Something like this: https://docs.databricks.com/applications/deep-learning/distributed-training/mnist-tensorflow-keras.html)

pandas_df = df.toPandas() # Not possible in real life
data_array = np.asarray(list(pandas_df.data.values))
data_array = data_array.reshape(data_array.shape[0], data_array.shape[1], -1, 1, order='F')
data_array = data_array.reshape(data_array.shape[0],data_array.shape[1],-1,1,1,order="F").transpose(0,1,3,2,-1)
# Some more numpy specific transformations ..

Here's an approach that didn't work:

@pandas_udf(ArrayType(IntegerType()), PandasUDFType.SCALAR)
def generate_feature(x):
    data_array = np.asarray(x)
    data_array = data_array.reshape(data_array.shape[0], ..
    ...
    return pd.Series(data_array)

df = df.withColumn("data_array", generate_feature(df.data))

Any updates? How do you solve this problem?

talentcat
– talentcat

2022-03-22 07:50:18 +00:00
Commented Mar 22, 2022 at 7:50 — talentcat
– talentcat, Commented Mar 22, 2022 at 7:50
@talentcat it's not possible sorry

maverik
– maverik

2022-04-13 20:07:33 +00:00
Commented Apr 13, 2022 at 20:07 — maverik
– maverik, Commented Apr 13, 2022 at 20:07

SriramN · Accepted Answer · 2020-03-06 11:30:56Z

0

I am trying to work on a similar case though using Images. I am looking towards Petastorm for doing this. You can save your data from Rdd to Parquet format and then use it in horovod.
- I am yet to test this.
- How to fetch the dataset in parts using ranks in horovod, needs to be tested too.
Just a tip that could help.
Thanks.

answered Mar 6, 2020 at 11:30

SriramN

5117 silver badges27 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Spark dataframe to numpy array via udf or without collecting to driver

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related