Schema overhaul

[pingsutw]

TL;DR

This PR introduces a new type and user experience for working with tabular data in Flyte. Please see the design doc for the complete background.

• The new type is the StructuredDataset. (see the IDL PR) We thought the name sufficiently generic and differentiated from FlyteSchema, which we're also keeping around for the time being (see more in the Migration section below). The main difference between the FlyteSchema and StructuredDataset types themselves is that columns in the latter are just named LiteralTypes. That is, columns are no longer constrained to be purely primitives.
• The contributor-experience has also changed substantially from FlyteSchema to something we hope is more straightforward and easier to understand, but we've tried to keep the user-experience as similar as possible.

Note:
The [kwtype(...)] syntax used for specifying columns and storage formats will be changing before the next non-beta release of flytekit, so if you try this out now, keep in mind you'll have some minor minor updates to make in a couple weeks.

Note:
To handle the new type, you will need to be on Propeller version v0.16.14 and Admin version v0.6.78 or later.

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

Design

Type Handling

We found the double-interaction between the core TypeEngine and the SchemaEngine too confusing.

In the new model, contributors will only be dealing with the StructuredDatasetTransformerEngine. You register your encoders/decoders with that only, and it in turn will handle interaction with the core TypeEngine.

If you have a new dataframe type that you want to be able to use in Flytekit, you'll just have to subclass StructuredDatasetEncoder and StructuredDatasetDecoder and implement the logic. Encoders are responsible for taking a dataframe in Python memory and turning it into a Flyte literal (replete with any uploading/writing of data) and decoders do the opposite (and should handle downloading of data).

Encoders and decoders are now distinct entities rather than two functions within one class because we felt the the asymmetry increases flexibility. One might choose to write a custom decoder for pandas but use the default encoder.

Encoder/Decoder (aka Handler) Selection

This PR ships with handlers to encode and decode pandas and spark dataframes. And users can of course add their own. How does the transformer engine pick the encoder/decoder?

The handlers in the new transformer-engine are keyed off of three things:

1. Python dataframe type
2. Storage protocol (s3, bq, gs, etc.)
3. Storage format (a string like "parquet", "csv", etc.)

Each en/decoder is required to have all three. When the transformer engine needs to convert a Python dataframe in memory to a Flyte literal, it will pick an encoder based on exact match of all three, except for the format. For the format, if an exact match cannot be found, it will then search for an encoder with the format "". The same pattern is used for decoders/when converting from a Flyte literal to a Python dataframe.

Default protocols and formats

When registering an en/decoder with the FLYTE_DATASET_TRANSFORMER singleton, you have the option of specifying whether the protocol and storage format that that handler handles, should be taken as the default format for that dataframe type.

For example, if you

from flytekit.types.structured.structured_dataset import FLYTE_DATASET_TRANSFORMER

FLYTE_DATASET_TRANSFORMER.register(MyOwnEncoder(my.DataFrame, "s3", "parquet"), default_for_type=True)
FLYTE_DATASET_TRANSFORMER.register(MyOtherEncoder(my.DataFrame, "gs", "avro"), default_for_type=False)

this will tell flytekit that whenever it sees

def t1() -> my.DataFrame:
    ...
    return my.DataFrame(...)

that it will assume that the format should be "parquet" and the protocol should be "s3", and it will look for the appropriate en/decoder.

Usage

Please see the tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py test for user examples. The basic rules are though that if you

def t1() -> some.DataFrame:
  return some.DataFrame(...)

then you'll get the default behavior (as determined by the default protocol and format). If you want to customize behavior (for example, specifying the uri that the data will be stored in), then you have to change the return signature and wrap the object returned inside the Python StructuredDataset class (as opposed to the literals.StructuredDataset model class).

def t1() -> StructuredDataset:
  return StructuredDataset(dataframe=some.DataFrame(...), uri="bq://some/path")

Reading is done much the same as before...

If you

def t1(a: pd.DataFrame): ...

then the default handler for pd.DataFrame will be used and whatever object the decoder returns will be what the task function is called with. Alternatively you can

def t1(a: StructuredDataset):
    a.open(pd.DataFrame).all()  # or
    # a.open(pd.DataFrame).iter  # if the decoder returns an iterator.

Migration and Compatibility

We'll be keeping the existing FlyteSchema object around for some time. However, there are still some backwards compatibility issues you may come across.

With this PR, pd.DataFrame and pyspark.sql.DataFrame will henceforth be stored in the Flyte backend as a StructuredDataset literal. Previously they were stored as Schema literals. We recommend that if you have any cached data frame tasks, that you bump the cache version and recompute.

This can also lead to compatibility issues between tasks. For example if you have:

@task
def t1() -> FlyteSchema: ...

@task
def t2(a: pd.DataFrame):  ...

where t2(a=t1()), once you upgrade to this version of flytekit, when t2 runs, it will try to use the new code, to read the old Schema literal. The reverse is also true

@task
def t1() -> pd.DataFrame: ...

@task
def t2(a: FlyteSchema):  ...

To handle this, we've added code to both the existing SchemaEngine and the new StructuredDataset engine to handle the other. Let us know if you have problems.

Plugins

We update the spark plugin made a few fixes to the great expectations plugin but most of the dataframe-related plugins remain untouched. We will be updating these in the weeks to come.

Upcoming Changes

There's a couple changes that we just felt would be better to do after we had a beta release to play with but they should really be thought of as part of this PR.

• The change to the column/format specifying syntax StructuredDataset[kwtypes(col_a=int), "my_parquet"]. This is going to be replaced within the next week or two.
• Column subsetting - if a task returns two columns but the next task only needs one of them, the dataframe provided should only have that one column.

Separately,

• The handlers shipped with this preserve the existing behavior of using aws/gcp clis to download and upload data. Faster handlers that use fsspec will be coming as well.
• Aforementioned plugin updates.

Detailed Description of Changes

Core changes

• Add override to TypeEngine.
• Strip Annotated in TypeEngine.

Compatibility

Plugins

• Moved Spark encoder and decoder into the spark plugin.
  • Changed encoder to use a randomly generated remote directory instead of a file.
• flytekit-sqlalchemy, flytekit-snowflake, flytekit-hive, flytekit-aws-athena need backend changes to handle the new StructuredDataset type in addition to FlyteSchema. Probably lower priority unless someone asks for it. After backend changes are done, we can update the Python type hints.
• flytekit-modin seems to operate on its own type, no interoperability present with the normal pandas.DataFrame type, so we don't have to worry about pandas -> modin compatibility.
• flytekit-dolt seems to be the same as modin.

Tracking Issues

User submitted, not the main issue: flyteorg/flyte#523
flyteorg/flyte#2074

Follow-up issue

NA