Skip to content

feat: add stratify param support to ml.model_selection.train_test_split method #815

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

Jump to bottom
Merged
merged 6 commits into from
Jul 10, 2024
Merged

feat: add stratify param support to ml.model_selection.train_test_split method #815

merged 6 commits into from
Jul 10, 2024

Conversation

GarrettWu
Copy link
Contributor

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

  • Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
  • Ensure the tests and linter pass
  • Code coverage does not decrease (if any source code was changed)
  • Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

@GarrettWu GarrettWu requested review from junyazhang and shobsi July 2, 2024 20:09
@GarrettWu GarrettWu requested review from a team as code owners July 2, 2024 20:09
@product-auto-label product-auto-label bot added the size: m Pull request size is medium. label Jul 2, 2024
@product-auto-label product-auto-label bot added the api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. label Jul 2, 2024
@GarrettWu GarrettWu assigned GarrettWu and unassigned orrbradford Jul 2, 2024
shobsi
shobsi approved these changes Jul 9, 2024
View reviewed changes
bigframes/ml/model_selection.py
test_dfs.append(test)

train_df = cast(
bpd.DataFrame, bpd.concat(train_dfs).drop(columns="bigframes_stratify_col")
Copy link
Contributor

Choose a reason for hiding this comment

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

concat complexity was much discussed in various threads (like RAG notebook, google next demo), e.g. https://screenshot.googleplex.com/7jehfHCAVJrmc3p. The number of unique values in the stratify col could be large to run into that. It would be a good idea to test and document where the limit lies.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes. Concat will grow BFET and SQL size. It may encounter SQL size or OOM errors when unique values size is too large. Added a note.