[Python] Enforcing Schema for PySpark Job

We have a bunch of disparate pyspark jobs and have a usecase where two jobs can be chained together to form a seperate 3rd job. So if you have one job 'A' and another job 'B' these jobs can be chained together to form another job 'C'. Obviously this can only happen if the output dataframe schema from 'A' is compatible with dataframe schema to 'B'. This is exactly what the dataset api achieves however unfortunately spark doesnt vend dataset for python, understandably because python is dynamically typed. However with the advent of typings we can achieve some sort of compile time safety and I can think of a couple of ways to achieve this, one through composition the other through inheritance.
Composition

from typing import TypeVar, Generic
from pyspark.sql import DataFrame
from enum import Enum

#Marker Enum
class Schema(Enum):
pass

T = TypeVar("T", bound = Schema)

class CustomDataSet(Generic[T]):
def __init__(self, dataframe: DataFrame) -> None:
self.dataframe = dataframe



Now I can use mypy and use composition like CustomDataSet[MySchema](dataframe). The problem out here is that MySchema is not dataframe object and this can be confusing to someone using this.

Inheritance

import abc
from enum import Enum

#Marker Enum
class Schema(Enum):
pass

class MyInterface(metaclass=abc.ABCMeta):

@abc.abstractmethod
def get_input_schema() -> Schema:
raise NotImplementedError

@abc.abstractmethod
def get_output_schema() -> Schema:
raise NotImplementedError



Now the two jobs 'A' and 'B' can implement the above interface. This seems a fairly verbose way of doing things and feel not a fan of Inheritance .

My question I have is there a more pythonic way of doing this preferably where I can use composition?

Continue reading...