Storing serialized objects in rows of dataframes

[kostjaigin]

Hi!

I am using your framework to apply a pre-trained torch model in pipeline. The question is more or less repeated here. I am trying serialize my graphs and to store their serializations in dataframe:

from pyspark.sql.types import BinaryType

classifier = ... # my pre-trained network network

data_frames = []

for test_graph in test_graphs:
    # transform subgraphs to pandas data frames:
    g = test_graph[0] # networkx object
    node_tag = test_graph[2] # list of integers
    df = nx.to_pandas_adjacency(g)
    df["tag"] = df.apply(lambda row: node_tag[row.name], axis=1)
    data_frames.append(pkl.dumps(df))

dataframe = spark.createDataFrame(data_frames, BinaryType())

model = create_spark_torch_model (
    classifier,
    inputCol='value',
    predictionCol='predictions'
)

predictions = model.transform(dataframe)
predictions.show(10) # call an action to perform transformations on data

The error I get comes from torch_distributed.py:

File "/Users/konstantinigin/opt/miniconda3/lib/python3.8/site-packages/sparktorch/torch_distributed.py", line 113, in predict_float
features = data.toArray().reshape((1, len(data)))
AttributeError: 'bytearray' object has no attribute 'toArray'

I took a look into the file and asked myself if I could fix it, but did not succeed. May be you have some idea what could be wrong here?