The answer is to use a custom Spark script with stratified sampling or the `train_test_split` function from scikit-learn with the `stratify` parameter. These two methods are correct because they both implement stratified splitting techniques for imbalanced datasets, ensuring that the proportion of each class in the target variable remains consistent across the training and testing splits. This is critical when working with skewed distributions, as random splitting could accidentally create a test set missing a rare class, leading to misleading model evaluation. On the AWS Certified Machine Learning Engineer Associate MLA-C01 exam, this question tests your understanding of data preparation in AWS Glue and SageMaker, often appearing as a scenario where you must choose between native Glue operations and custom code. A common trap is selecting a simple random split, which ignores class imbalance. Remember the memory tip: "Stratify saves the minority" — always look for the `stratify` parameter or a custom sampling script when preserving target distribution is required.
MLA-C01 Data Preparation for Machine Learning Practice Question
This MLA-C01 practice question tests your understanding of data preparation for machine learning. Match the stated requirement to the specific cloud service, access model, or configuration option — many options are valid in isolation but not for this scenario. After answering, compare your reasoning against the explanation and wrong-answer breakdown below. Once you have made your selection, read the full explanation to reinforce the concept and understand why each distractor is designed to mislead on exam day.
A data engineer is using AWS Glue to prepare a dataset for ML. The engineer wants to split the dataset into training and testing sets while preserving the distribution of the target variable. Which TWO methods achieve this goal? (Select TWO)
Answer the question above first, then reveal the full breakdown to understand why each option is right or wrong.
Correct answer & explanation
✓
Use the `train_test_split` function from scikit-learn in a SageMaker notebook
Option B is correct because the `train_test_split` function from scikit-learn supports the `stratify` parameter, which preserves the distribution of the target variable when splitting a dataset into training and testing sets. This is a standard, reliable method for stratified splitting in Python-based ML workflows, and it can be used directly in a SageMaker notebook.
Key principle: Answer the scenario, not the keyword: identify the specific constraint before choosing the most familiar-sounding option.
Answer analysis
Option-by-option breakdown
For each option: why learners choose it and why it is or isn't the right answer here.
✗
Use Amazon Athena to create views with random sampling
Why it's wrong here
Athena random sampling does not guarantee stratified splits.
✓
Use the `train_test_split` function from scikit-learn in a SageMaker notebook
Why this is correct
The stratify parameter maintains class proportions.
Related concept
Read the scenario before looking for a memorised answer.
✗
Use AWS Glue's built-in random split transform
Why it's wrong here
Random split does not preserve target distribution.
✓
Use a custom Spark script with stratified sampling
Why this is correct
Stratified sampling ensures proportional representation of target classes.
Related concept
Read the scenario before looking for a memorised answer.
✗
Use Amazon SageMaker's built-in SplitType parameter in a Processing Job
Why it's wrong here
SplitType only supports random division.
Common exam traps
Common exam trap: answer the scenario, not the keyword
The trap here is that candidates often confuse random splitting (which is available in many tools like Glue and Athena) with stratified splitting, assuming that any 'random' operation preserves distribution, but only stratified methods explicitly maintain class proportions.
Detailed technical explanation
How to think about this question
Stratified splitting ensures that each class or category of the target variable is proportionally represented in both training and testing sets, which is critical for imbalanced datasets to avoid biased model evaluation. Under the hood, scikit-learn's `train_test_split` with `stratify=y` performs a group-wise shuffle and split based on the class labels, using the `StratifiedShuffleSplit` algorithm. In a real-world scenario, if you have a binary classification dataset with 90% negative and 10% positive samples, a simple random split could accidentally place all positive samples in the test set, but stratified splitting maintains the 90/10 ratio in both splits.
KKey Concepts to Remember
Read the scenario before looking for a memorised answer.
Find the constraint that changes the correct option.
Eliminate answers that are true in general but not in this case.
TExam Day Tips
→Watch for words such as best, first, most likely and least administrative effort.
→Review why wrong options are wrong, not only why the correct option is correct.
Key takeaway
Answer the scenario, not the keyword: identify the specific constraint before choosing the most familiar-sounding option.
Real-world example
How this comes up in practice
A cloud solutions architect for a retail company is evaluating services for a new workload. The correct answer here reflects best practice for the specific scenario described — not a general cloud recommendation. Answer the scenario, not the keyword: identify the specific constraint before choosing the most familiar-sounding option. Cloud exam questions reward reading the constraint carefully: the same technology can be right or wrong depending on the use case.
What to study next
Got this wrong? Here's your next step.
Identify which exam domain this question belongs to, review the core concept, then practise similar questions from the same domain.
Data Preparation for Machine Learning — This question tests Data Preparation for Machine Learning — Read the scenario before looking for a memorised answer..
What is the correct answer to this question?
The correct answer is: Use the `train_test_split` function from scikit-learn in a SageMaker notebook — Option B is correct because the `train_test_split` function from scikit-learn supports the `stratify` parameter, which preserves the distribution of the target variable when splitting a dataset into training and testing sets. This is a standard, reliable method for stratified splitting in Python-based ML workflows, and it can be used directly in a SageMaker notebook.
What should I do if I get this MLA-C01 question wrong?
Identify which exam domain this question belongs to, review the core concept, then practise similar questions from the same domain.
What is the key concept behind this question?
Read the scenario before looking for a memorised answer.
About these practice questions
Courseiva creates original exam-style practice questions with explanations and wrong-answer analysis. It does not publish real exam questions, exam dumps, or protected exam content. Learn why practice questions differ from exam dumps →
These questions test the same concept from different angles. Work through them to make sure you can recognise it however the exam phrases it.
Variation 1. A data scientist needs to split a dataset into training, validation, and test sets. The dataset has a categorical target variable with imbalanced class distribution. Which splitting technique ensures that each subset has a similar proportion of each class?
medium
A.K-fold cross-validation split
B.Chronological split
✓ C.Stratified split
D.Random split
Why C: Option C is correct because stratified splitting preserves the original class proportions in each subset (training, validation, test) by sampling each class independently. This is critical for imbalanced datasets to avoid skewed distributions that could bias model evaluation or training.
Last reviewed: Jun 24, 2026
Question Discussion
Share a tip, memory trick, or ask about the reasoning behind this question. Do not post real exam questions, leaked content, braindumps, or copyrighted exam material. Comments are moderated and may be removed without notice.
This MLA-C01 practice question is part of Courseiva's free Amazon Web Services certification practice question bank. Courseiva provides original exam-style practice questions with explanations, topic-based practice, mock exams, readiness tracking, and study analytics to help learners prepare for the MLA-C01 exam.
Question Discussion
Share a tip, memory trick, or ask about the reasoning behind this question. Do not post real exam questions, leaked content, braindumps, or copyrighted exam material. Comments are moderated and may be removed without notice.
Sign in to join the discussion.