How to retain attributes of inherited Spark DataFrame Class following a Spark operation on that class

Viewed 55

I create a new class called NewDataFrame with attribute a_string:

import numpy as np
import pandas as pd
from pyspark.sql import DataFrame

class NewDataFrame(DataFrame):

    def __init__(self, df):
        super().__init__(df._jdf,df.sql_ctx)
        self.a_string = "Hello, World."

I use the class on some data and am able to print out a_string:

data = {
     'a': ['yellow', 'red']
    ,'b': [1, 2]
}

df = pd.DataFrame(data)
sdf = spark.createDataFrame(df)

temp = NewDataFrame(sdf)
temp.a_string

Out[]: Hello, World.

Now, I filter temp to a subset and try to output a_string and receive an error because the filter method returns a DataFrame, not NewDataFrame.

temp = temp.filter("a='yellow'")
temp.a_string

Out[]: 'DataFrame' object has no attribute 'a_string'

To keep the attribute in the result of a filter, I have tried creating a new method on the NewDataFrame class which performs the filter and then feeds the result back into a NewDataFrame class, which works, but I do not want to rewrite all the Spark functions in this manner.

Is there a way for the class to have access to the full range of DataFrame methods while still retaining the attributes I define in NewDataFrame?

0 Answers
Related