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?