Skip to content

fix: ml.model_selection.train_test_split index to match in unordered mode #2283

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GarrettWu wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix: ml.model_selection.train_test_split index to match in unordered mode #2283

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2b5f122
fix: ml.model_selection.train_test_split index to match in unordered …
GarrettWu Nov 19, 2025
9e2cd71
fix
GarrettWu Nov 20, 2025
b29c6f3
move cache()
GarrettWu Nov 24, 2025
b2815b0
Merge remote-tracking branch 'github/main' into garrettwu-split
GarrettWu Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions bigframes/ml/model_selection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ def _stratify_split(df: bpd.DataFrame, stratify: bpd.Series) -> List[bpd.DataFra
joined_df = dfs[0]
for df in dfs[1:]:
joined_df = joined_df.join(df, how="outer")
joined_df = joined_df.cache()
if stratify is None:
joined_df_train, joined_df_test = joined_df._split(
fracs=(train_size, test_size), random_state=random_state
Expand All @@ -120,8 +121,8 @@ def _stratify_split(df: bpd.DataFrame, stratify: bpd.Series) -> List[bpd.DataFra
results = []
for array in arrays:
columns = array.name if isinstance(array, bpd.Series) else array.columns
results.append(joined_df_train[columns])
results.append(joined_df_test[columns])
results.append(joined_df_train[columns].cache())
results.append(joined_df_test[columns].cache())
Copy link
Contributor

@TrevorBergeron TrevorBergeron Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a lot of .cache() calls. I think where the caching ideally happens is actually inside the block.split method. This way, the ordering is locked in, but only a single table is cached total, which should be a lot faster.

Copy link
Contributor Author

@GarrettWu GarrettWu Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can put the caches outside of the loop, which removes some queries. But otherwise (caching anywhere in block.split) it doesn't work. Do you have an insight why is it?

Copy link
Contributor

@TrevorBergeron TrevorBergeron Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm, really? would expect caching anywhere around this area:

python-bigquery-dataframes/bigframes/core/blocks.py

Lines 901 to 913 in b487cf1

block, string_ordering_col = block.apply_unary_op(
ordering_col, ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE)
)
# Apply hash method to sum col and order by it.
block, string_sum_col = block.apply_binary_op(
string_ordering_col, random_state_col, ops.strconcat_op
)
block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op)
block = block.order_by(
[ordering.OrderingExpression(ex.deref(hash_string_sum_col))]
)
would work ok (ideally at the end of this block).

Copy link
Contributor Author

@GarrettWu GarrettWu Nov 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, no matter where, within the block.split, it won't work. Only do a cache() to the end results would help.

screen/6A2RFRNf9m96Qvo

Would it be a bug in some deeper code?


return results

Expand Down
18 changes: 18 additions & 0 deletions tests/system/small/ml/test_model_selection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,24 @@ def test_train_test_split_default_correct_shape(df_fixture, request):
assert y_test.shape == (86, 1)


def test_train_test_split_default_unordered_same_index(
unordered_session, penguins_pandas_df_default_index
):
df = unordered_session.read_pandas(penguins_pandas_df_default_index)
X = df[
[
"species",
"island",
"culmen_length_mm",
]
]
y = df[["body_mass_g"]]
X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y)

pd.testing.assert_index_equal(X_train.to_pandas().index, y_train.to_pandas().index)
pd.testing.assert_index_equal(X_test.to_pandas().index, y_test.to_pandas().index)


def test_train_test_split_series_default_correct_shape(penguins_df_default_index):
X = penguins_df_default_index[["species"]]
y = penguins_df_default_index["body_mass_g"]
Expand Down