Test a generic starscream stage

sample_code.py

from __future__ import unicode_literals

import pytest

from starscream.pipeline.stage import TransformStage

from pyspark.sql import functions as F, types as T



from starscream.contract import Contract
from starscream.utils.dataframe import as_dicts, from_dicts

import pyspark.sql.types as T

contract = Contract({
    'foo': {'type': T.LongType()},
    'bar': {'type': T.LongType()},
})

input_df = from_dicts(sc, contract, [
    {'foo': 1, 'bar': 2},
    {'foo': 2, 'bar': 3},
    {'foo': 3, 'bar': 4},
    {'foo': 4, 'bar': 5},
    {'foo': 5, 'bar': 6},
])

input_df.select(['foo', 'bar']).show()

class MyStage(TransformStage):
    
    OUTPUT = Contract({
        'foo': {'type': T.LongType()},
        'bar': {'type': T.LongType()},
        'baz': {'type': T.LongType()},
    })
    
    def apply(self, sc, my_input_df):
        return (
            my_input_df
            .withColumn('baz', F.col('foo') + F.col('bar'))
        )

output_df = MyStage().apply(sc, input_df)
output_df.select(['foo', 'bar', 'baz']).show()