fix(utils): remove redundant column in train_test_split() (ibis-project#131)

Author: jitingxu1

ibis_ml/utils/_split.py

@@ -90,6 +90,6 @@ def train_test_split(
         < int((1 - test_size) * num_buckets)
     )
 
-    return table[table.train].drop(["combined_key"]), table[~table.train].drop(
-        ["combined_key"]
+    return table[table.train].drop(["combined_key", "train"]), table[~table.train].drop(
+        ["combined_key", "train"]
     )

tests/test_train_test_split.py

@@ -16,6 +16,8 @@ def test_train_test_split():
     # Check counts and overlaps in train and test dataset
     assert train_table.count().execute() + test_table.count().execute() == N
     assert train_table.intersect(test_table).count().execute() == 0
+    assert set(train_table.columns) == set(table.columns)
+    assert set(test_table.columns) == set(table.columns)
 
     # Check reproducibility
     reproduced_train_table, reproduced_test_table = ml.train_test_split(