Thoughts on a Flow step function decorator that injects a variable as a function parameter?

class HelloFlow(FlowSpec):
    @datasets.dataset("my_dataset")
    @step
    def start(self, my_dataset: Dataset):
       pass

This would change how Metaflow parses a Flow

join

step https://github.com/zillow/metaflow/blob/tz/datasets/metaflow/graph.py#L114-L116

elif len(self.out_funcs) == 1:
                    # TYPE: linear
                    if self.num_args > 1:
                        self.type = 'join'
                    else:
                        self.type = 'linear

The conundrum for me is that the definition of a

join

step is a Flow function with

> 1

argument...

right.. that's the main conundrum

we could, but because of scoping that may be confusing

you don't want the artifact to be visible outside this step?

one (an Applied Scientist user, or any engineer) would likely think that a function decorator introduces a variable at function scope

this small hack works

extra

is a tuple 🙂

that's true

I wish a join step had to marked or have variable named

inputs

(making it a keyword)

@step
def join_step(self, inputs):
   pass

Do all join steps use

inputs

as variable name? I bet you can't guarantee that...

right. Given that we don't want to break existing flows, we'd just need to find a way to make the change in a backwards-compatible manner. I am optimistic that it is doable but requires some thinking

it's what we currently do 🙂

class VanillaDatasetFlow(FlowSpec):
    @datasets.dataset(name="vanilla_dataset", partition_by="group", mode=Mode.Write)
    @step
    def start(self):
        df = pd.DataFrame({"group": ["A", "A", "A", "B", "B", "B"], "value": [1, 2, 3, 4, 5, 6]})
        print(f"saving: {df=}")

        # Example of writing to a vanilla dataset 
        self.vanilla_dataset.write(df)

        self.next(self.end)

    @step
    def end(self):
        print(f"I have dataset \n{self.vanilla_dataset=}")

        # vanilla_dataset read_pandas()
        print("self.vanilla_dataset.read_pandas:\n", self.vanilla_dataset.read_pandas())

        # vanilla_dataset read_spark()
        with SparkManager(auth=False):
            spark_df: DataFrame = self.vanilla_dataset.read_spark()
            spark_df.show()

        # vanilla_dataset read_dask()
        dask_df: DataFrame = self.vanilla_dataset.read_dask()
        dask_df = dask_df[dask_df.value < 4]
        print("self.vanilla_dataset.read_dask: [value < 4]\n", dask_df.compute())

cool, that should work