Add a new shape field to ColumnSchema

[karlhigley]

This creates a place to store shape information for all dimensions of the data across both array/tensor and dataframe formats. In contrast to the existing "value_count" property (which only records the value counts of the lists in list field, this attribute is intended to capture the size of all dimensions of the data (the batch dimension, the list lengths, embedding sizes, etc.)

A few significant design elements:

• Defining a Shape dataclass allows us to make them immutable, which we wouldn't get if we sub-classed tuple like some other frameworks do. For compatibility with those other frameworks (and just general ease of use), the Shape constructor accepts a tuple of dimensions, so it should be relatively straightforward to (for example) create a Merlin shape from a Tensorflow shape.
• Shapes are stored in dtypes, since shape information is how we know a particular column is a list, and several frameworks we use (e.g. cuDF, Feast) have either a general List dtype or specific list dtypes like int32_list for each element type. The API has been designed to hide this from calling code as much as possible, but it does influence how shapes work to some extent.
• The is_list and is_ragged flags are now computed from the shape and can no longer be set by providing them as constructor args. We tried to find a way to maintain that, but it's not possible to use dataclass init vars with the same name as a defined property. The best workaround for this we could come up with was to allow is_list and is_ragged to be set from the properties dictionary, in which case we do our best to infer a shape from the values of these flags.

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/core/review/pr-195

[EvenOldridge]

I like the level of detail into the error messaging and the comprehensive tests. Assuming the merge of shapes is meant to be future work this looks good.

[karlhigley]

@oliverholworthy I updated this PR to keep is_list and is_ragged as properties that are set in the __post_init__ to match the shape, which really does clean things up quite a bit. I think it still ends up being a breaking change that requires coordination with the other repos though, because validating that the provided flags are consistent with the shape turns up places where that isn't currently true.

PR changed significantly in response to feedback after approval

[oliverholworthy]

I think it still ends up being a breaking change that requires coordination with the other repos though, because validating that the provided flags are consistent with the shape turns up places where that isn't currently true.

Looks like we'll be setting a requirement for the shape to be specified if is_ragged=False is specified, the fixed dimension will be required to be passed. Would it be possible or helpful with the migration to loosen that restriction further and only check the shape if it's used explictly (either with the dims passed to the ColumnSchema constructor or with_shape)

[karlhigley]

Would it be possible or helpful with the migration to loosen that restriction further and only check the shape if it's used explictly?

I shifted the validation in this general direction, so that only information that's explicitly provided gets validated. I'm not sure about lifting the restriction that some kind of shape info has to be provided (through the dtype's shape, the dims, or the value counts) when is_ragged=False because that will put us in a position where we sometimes don't have a shape. That may complicate the downstream changes as we try to migrate over. 🤔

[oliverholworthy]

because that will put us in a position where we sometimes don't have a shape. That may complicate the downstream changes as we try to migrate over

I suppose we're already in the position where we sometimes don't have a shape, since there's nothing in here currently to enforce having a value count when is_ragged=False. Adding this requirement in this PR forces us to address that now. I support that choice if we're commited to making the changes to the downstream libs right now. If it's mostly a question of how we'll identify which parts of the code need to change, we could achieve that with an optional strict mode or something along those lines.

And closely related to this: to what extent will this change make to loading previously saved schemas (e.g. from an NVTabular workflow and saved schema file). Will we be able to load schemas saved with a prior version after merging this? We'd presumably prefer to be backwards compatible if possible, however if it's a necessary breaking change across the saved file format, then we might need some extra docs about this when we release the next version.

[karlhigley]

Adding this requirement in this PR forces us to address that now. I support that choice if we're commited to making the changes to the downstream libs right now. If it's mostly a question of how we'll identify which parts of the code need to change, we could achieve that with an optional strict mode or something along those lines.

We're thinking along similar lines then; we've disabled the is_ragged=False without shape info validation for now, and will use it to identify the places where changes are needed in further dev work. 👍🏻

to what extent will this change make to loading previously saved schemas (e.g. from an NVTabular workflow and saved schema file). Will we be able to load schemas saved with a prior version after merging this? We'd presumably prefer to be backwards compatible if possible, however if it's a necessary breaking change across the saved file format, then we might need some extra docs about this when we release the next version.

I believe the change is backward but maybe not forward compatible, since the allowed values of value counts are slightly different in a more permissive direction after this change.

[oliverholworthy]

Looks like this is in good shape! ✨