Skip to content

Commit

Permalink
[SPARK-38271] PoissonSampler may output more rows than MaxRows
Browse files Browse the repository at this point in the history
### What changes were proposed in this pull request?
when `replacement=true`, `Sample.maxRows` returns `None`

### Why are the changes needed?
the underlying impl of `SampleExec` can not guarantee that its number of output rows <= `Sample.maxRows`

```
scala> val df = spark.range(0, 1000)
df: org.apache.spark.sql.Dataset[Long] = [id: bigint]

scala> df.count
res0: Long = 1000

scala> df.sample(true, 0.999999, 10).count
res1: Long = 1004
```

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
existing testsuites

Closes #35593 from zhengruifeng/fix_sample_maxRows.

Authored-by: Ruifeng Zheng <ruifengz@foxmail.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit b683279)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
  • Loading branch information
zhengruifeng authored and cloud-fan committed Feb 22, 2022
1 parent 79099cf commit 7d36329
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -1344,7 +1344,11 @@ case class Sample(
s"Sampling fraction ($fraction) must be on interval [0, 1] without replacement")
}

override def maxRows: Option[Long] = child.maxRows
override def maxRows: Option[Long] = {
// when withReplacement is true, PoissonSampler is applied in SampleExec,
// which may output more rows than child.maxRows.
if (withReplacement) None else child.maxRows
}
override def output: Seq[Attribute] = child.output

override protected def withNewChildInternal(newChild: LogicalPlan): Sample =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,19 @@ class CombiningLimitsSuite extends PlanTest {
)
}

test("SPARK-38271: PoissonSampler may output more rows than child.maxRows") {
val query = testRelation.select().sample(0, 0.2, true, 1)
assert(query.maxRows.isEmpty)
val optimized = Optimize.execute(query.analyze)
assert(optimized.maxRows.isEmpty)
// can not eliminate Limit since Sample.maxRows is None
checkPlanAndMaxRow(
query.limit(10),
query.limit(10),
10
)
}

test("SPARK-33497: Eliminate Limit if Deduplicate max rows not larger than Limit") {
checkPlanAndMaxRow(
testRelation.deduplicate("a".attr).limit(10),
Expand Down

0 comments on commit 7d36329

Please sign in to comment.