-
Notifications
You must be signed in to change notification settings - Fork 174
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
3cb7252
commit 9d76be3
Showing
3 changed files
with
62 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
use eyre::{bail, WrapErr}; | ||
use tracing::warn; | ||
|
||
use crate::translation::{to_logical_plan, Plan}; | ||
|
||
pub fn sample(sample: spark_connect::Sample) -> eyre::Result<Plan> { | ||
let spark_connect::Sample { | ||
input, | ||
lower_bound, | ||
upper_bound, | ||
with_replacement, | ||
seed, | ||
deterministic_order, | ||
} = sample; | ||
|
||
let Some(input) = input else { | ||
bail!("Input is required"); | ||
}; | ||
|
||
let mut plan = to_logical_plan(*input)?; | ||
|
||
// Calculate fraction from bounds | ||
// todo: is this correct? | ||
let fraction = upper_bound - lower_bound; | ||
|
||
let with_replacement = with_replacement.unwrap_or(false); | ||
|
||
// we do not care about sign change | ||
let seed = seed.map(|seed| seed as u64); | ||
|
||
if deterministic_order { | ||
warn!("Deterministic order is not yet supported"); | ||
} | ||
|
||
plan.builder = plan | ||
.builder | ||
.sample(fraction, with_replacement, seed) | ||
.wrap_err("Failed to apply sample to logical plan")?; | ||
|
||
Ok(plan) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
from __future__ import annotations | ||
|
||
|
||
def test_sample(spark_session): | ||
# Create a range DataFrame | ||
df = spark_session.range(100) | ||
|
||
# Test sample with fraction | ||
sampled_df = df.sample(fraction=0.1, seed=42) | ||
sampled_rows = sampled_df.collect() | ||
|
||
# Verify sample size is roughly 10% of original | ||
sample_size = len(sampled_rows) | ||
assert 5 <= sample_size <= 15, f"Sample size {sample_size} should be roughly 10 rows" | ||
|
||
# Verify sampled values are within original range | ||
for row in sampled_rows: | ||
assert 0 <= row["id"] < 100, f"Sampled value {row['id']} outside valid range" |