How to use train_test_split

[wiso]

train_test_split from sklearn.model_selection is very popular, but it doesn't give what I expect. Example

data_train, data_test = train_test_split(data_all, test_size=0.2, random_state=42)
data_train

now data_train is a List[awkward.highlevel.Record]

[<Record, ...>, <Record, ...]

instead, I would have expected a awkward.highlevel.Array

I can use ak.merge_union_of_records(data_train, axis=0) but it is extremely slow.

[agoose77]

Are you invoking test_train_split on Awkward Array data?

[wiso]

yes, sure