An example feature class

pyspark_feature_a_to_b_ratio_example.py

from pyspark.sql import functions as F


class FeatureAToBRatio(object):
    feature_name = 'a_to_b_ratio'
    default_value = 0.

    def calculate(self, df):
        """
        Given a dataframe that contains columns a and b,
        calculate a to b ratio. If b is 0 then the result will be
        the feature's default value.
        """
        df = df.withColumn(
            self.feature_name,
            F.when(
                F.col('b') > 0.,
                (
                    F.col('a').cast('float') / F.col('b').cast('float')
                )
            ).otherwise(
                self.default_value
            )
        ).fillna({self.feature_name: self.default_value})

        return df