Union Types #minor

[maximsmol]

TL;DR

Tagged Union types to support typing.Union in Python

Inherited from #230

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

https://github.com/maximsmol/flyte/blob/master/rfc/core%20language/sum-types.md

Tracking Issue

flyteorg/flyte#1349

Follow-up issue

NA

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[EngHabu]

Should we replace this with:

Suggested change

	UnionType type = 2;
	uint64 tag = 3;
	LiteralType type = 2;

Or in other words, what's the benefit of doing this vs just stating the LiteralType?

cc @wild-endeavor

[EngHabu]

@maximsmol after discussing some more this with the team, I think there are more changes needed in this PR... to explain that let's look at how this is all wired... but let me first explain what I think is needed first:

1. It's probably enough to assign the index of the type instead of carrying the entirety of the UnionType with us.
2. We need to modify this to:

// An input/output binding of a variable to either static value or a node output.
message Binding {
    // Variable name must match an input/output variable of the node.
    string var = 1;

    // Data to use to bind this variable.
    BindingData binding = 2;

    // UniontTypeInfo encapsulates extra binding information in case the type of `var` is a Union Type.
    UnionTypeInfo uniont_type_info = 3;
}

// UnionTypeInfo encapsulates extra binding information in case the type of `var` is a Union Type.
message UnionTypeInfo {
    // Maybe pick a better name
    int type_index = 1;
}

3. In the compiler, when we validate/build the bindings, we can fill in that field...

4. In propeller, at runtime, we can check if the binding has that UnionTypeInfo set, then we can create a wrapper UnionType literal with the type_index populated

Now to step back for a sec and discuss why are we doing all of this:
Flyte Literals are lossy (so are all "values" in languages), languages rely on storing type information along side the values to be able to do type enforcement/casting and other checks... Flyte does not require that from SDKs at runtime because we do all the validation at compile time when we actually have access to all the type information we need.
During the compilation is when we can validate the bindings built by the SDK to determine whether outputs match inputs. That's when we can also strongly bind the Union type that will be produced.

There is a caveat to this which is what happens if someone defines the output of a task as a union... I feel like I'm ok rejecting that for now... the compiler can validate and fail this case since 1) Not sure if it's a good UX to begin with and 2) will make it harder to do this binding at compile time...

[maximsmol]

1. Agreed, was originally there to deal with Unions with different orderings in Python (e.g. Union[str, int] and Union[int, str] would be incompatible if all we had was an integer index), but given the switch to string tags this is no longer needed

2, 3, 4. It's a good question whether we should allow using a value as a union implicitly as we would have to guess which instantiation was meant and, in languages like C++, the choice would be arbitrary. For example, how do we instantiate std::variant<int, int> i.e. Tag0 Int | Tag1 Int with a LiteralRepr containing an integer? Do we pick the first or the second variant? This is similarly the case for languages with algebraic data types (symbolically tagged unions). For Python we also have to make an arbitrary choice between two type transformers when instantiating Union[MyInt | int] but this case could be solved with somehow passing along the return type's transformer's name from the Python source annotation. This trick would not work for std::variant

In any case, if we decided to somehow do this automatic wrapping, we already have the UnionType accessible through validateBinding's expectedType argument, so I don't see a reason to pass it along. I might be misunderstanding something though.

[maximsmol]

Thought about it some more, and I think we actually do want to do automatic wrapping iff unambiguous

Example:

def f() -> str:
  return "10"
  
def g(x: Union[int, str]) -> str:
  return str(x)
  
def h(x: Union[MyStr, str]) -> str:
  return str(x)
  
g(f()) # all good, unambiguous
h(f()) # compilation failure

def workaround(x: str) -> Union[MyStr, str]:
  return MyStr(x)
  
h(workaround(f())) # all good, unambiguous (union literal includes disambiguating tag)

This matches the C++ std::variant API and produces reasonable behavior in Python, in Haskell the SDK could provide an extra check

[maximsmol]

closed by mistake (automatically by github after branch rename)

[EngHabu]

A huge thank you to @maximsmol and @eapolinario for pushing through this.

[EngHabu]

somethings with tabs vs spaces...

[welcome]

Congrats on merging your first pull request! 🎉