Design Draft

[wenleix]

[wenleix]

• HomoNumericDataFrame: Useful to model a group of dense features (backed by 2D Tensor with shape (batch, num_features) )
• GenericDataFrame: Backed by Dict[str, Union[ColumnBase, DataFrameBase]]
• GenericListColumn: The classic "flattened value" + offset encoding:

  class GenericListColumn
  values: Union[ColumnBase, DataFrameBase]
  lengths: Optional[torch.Tensor]
  offsets: Optional[torch.Tensor]

  Note in standard columnar database / arrow encoding, we should provide offsets. But in ML seems lengths are often more preferred :) . (Also, Velox also use both offsets and lengths )

• NumericListColumn: To represent List[Numeric] column; seems no difference from GenericListColumn, except values can be a 1D Tensor directly (e.g. avoid one level of wrapper; but not sure whether it's a big deal :) )
• Unless it's List[List[Numeric]], we can leverage NestedTensor as the internal representation?

Other possible useful specialized column

• A special column to encode Struct[NumericList], simliar to KeyedJaggedTensor. Will this be useful in preproc time???

Open question:

• Should DataFrameBase be a subclass of ColumnBase? -- or it can just be an type alias to StructColumnBase? :)

References:

• The early "demo" backend prototype from TorchArrow: https://github.com/pytorch/torcharrow/tree/9927a9c428c99ec8c4be3352c318ed2c66d36990/torcharrow/demo_rt
• Trino/Presto's Arrow in-memory format implementation on Java, e.g. RowBlock, MapBlock
• Velox's Arrow in-memory format implementation, e.g. RowVector, ArrayVector, MapVector

[wenleix]

Revised design based on initial iterations:

• Open to NumericListColumn specialization (which can support not only List[Numeric] but also List[List*[Numeric]] so the underlying storage can potentially leverage NestedTensor)
• For group dense features, FixedSizeListCol seems to be a quite good fit for feature-first Tensor layout.
• FixedTypeStructColumn seems to capture the use case of KeyedJaggedTensor quite well :)
• For now just have StructColumn while it provides some DataFrame APIs , to reduce the concepts introduced to user, as Axolotls is designed to be a model developer directly interfacing library.