Skip to content

[SPARK-50520][PySpark] Respect timeout in df.rdd.countApprox() #56060

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

Closed
rishav23 wants to merge 2 commits into from
Closed

[SPARK-50520][PySpark] Respect timeout in df.rdd.countApprox() #56060

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cb063ef
[SPARK-50520][PySpark] Respect timeout semantics in approximate RDD a…
rishav23 May 22, 2026
0a46c2b
[SPARK-50520][PySpark] Trigger CI
rishav23 May 22, 2026
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
6 changes: 4 additions & 2 deletions python/pyspark/core/rdd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4843,7 +4843,8 @@ def sumApprox(
jrdd = self.mapPartitions(lambda it: [float(sum(it))])._to_java_object_rdd()
assert self.ctx._jvm is not None
jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
r = jdrdd.sumApprox(timeout, confidence).getFinalValue()
partial = jdrdd.sumApprox(timeout, confidence)
r = partial.initialValue()

@cloud-fan cloud-fan Jun 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The fix here is correct, but it's incomplete: meanApprox (line 4885) is the direct peer of sumApprox -- same JavaDoubleRDD approximate action returning a PartialResult[BoundedDouble] -- yet it still calls .getFinalValue(), which blocks until the entire job finishes (PartialResult.getFinalValue waits until setFinalValue, fired only when all tasks complete). So meanApprox(timeout=...) retains the exact bug this PR fixes for sumApprox/countApprox, and the two parallel paths now diverge.

The PR description frames the problem generically ("PySpark approximate RDD actions currently call getFinalValue() ..."), which reads as covering all three actions. Suggest applying the same change to meanApprox:

Suggested change
r = partial.initialValue()
partial = jdrdd.meanApprox(timeout, confidence)
r = partial.initialValue()

(this suggestion is for line 4885, shown here for reference). If meanApprox is intentionally out of scope, please narrow the PR description to say so explicitly.

@rishav23 rishav23 Jun 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks for the review.
I've updated meanApprox to use the same timeout-aware handling as sumApprox by switching from getFinalValue() to initialValue().
I also updated the regression test to use a deliberately slow workload so that the timeout behavior is observable. The previous version could complete quickly even with the blocking implementation and therefore did not reliably detect the regression.
Additionally, the test now cleans up the background approximate job before proceeding to avoid interference with subsequent tests.
Verified with: python/run-tests.py --testnames pyspark.tests.test_rdd

return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high())

def meanApprox(
Expand Down Expand Up @@ -4881,7 +4882,8 @@ def meanApprox(
jrdd = self.map(float)._to_java_object_rdd()
assert self.ctx._jvm is not None
jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
r = jdrdd.meanApprox(timeout, confidence).getFinalValue()
partial = jdrdd.meanApprox(timeout, confidence)

@cloud-fan cloud-fan Jun 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Non-blocking: this meanApprox timeout fix isn't covered by a regression test. test_count_approx_respects_timeout exercises countApprox, which delegates to sumApprox -- so it guards the sumApprox path, but meanApprox makes a separate jdrdd.meanApprox(...) call and is the exact path that was missed in round 1. A small test mirroring the count one (slow workload, short timeout, assert elapsed is well under full completion) would prevent a silent re-regression on meanApprox alone.

r = partial.initialValue()
return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high())

def countApproxDistinct(self: "RDD[T]", relativeSD: float = 0.05) -> int:
Expand Down
30 changes: 30 additions & 0 deletions python/pyspark/tests/test_rdd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -638,6 +638,36 @@ def test_distinct(self):
self.assertEqual(result.getNumPartitions(), 5)
self.assertEqual(result.count(), 3)

def test_count_approx_respects_timeout(self):
def slow(x):
time.sleep(1)
return x

rdd = self.sc.parallelize(range(20), 20).map(slow)
start = time.time()
rdd.countApprox(timeout=100)
elapsed = time.time() - start
self.assertLess(elapsed, 2)
# Cancel the background approximate job before subsequent tests run.
self.sc.cancelAllJobs()

def test_count_approx_returns_exact_when_completed(self):
rdd = self.sc.parallelize(range(1000), 8)
self.assertEqual(rdd.countApprox(timeout=5000), 1000)

def test_mean_approx_respects_timeout(self):
def slow(x):
time.sleep(1)
return float(x)

rdd = self.sc.parallelize(range(20), 20).map(slow)
start = time.time()
rdd.meanApprox(timeout=100)
elapsed = time.time() - start
self.assertLess(elapsed, 2)
# Cancel the background approximate job before subsequent tests run.
self.sc.cancelAllJobs()

def test_external_group_by_key(self):
self.sc._conf.set("spark.python.worker.memory", "1m")
N = 2000001
Expand Down