[SPARK-52809][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric update by viirya · Pull Request #51503 · apache/spark

viirya · 2025-07-15T20:11:48Z

What changes were proposed in this pull request?

This patch adds only one task completion listener for metric updating for DataSourceRDD, instead of adding separate one for each partition iterator.

Why are the changes needed?

For each partition iterator, currently we add one task completion listener used to update final metrics if the task is stopped early. In the listener, the reader and iterator are held. So if the partition is normally exhausted, the references cannot be released early. It is a problem especially if the references are heavy as reported by apache/iceberg#13297.

Since the purpose of the callback is to update the final metrics if the task is stopped early, it means that we only need to do it for the last partition iterator. So we don't need set up a listener for each partition iterator. Thus, we can set up just one listener for all iterator. Once we advance to next partition iterator, we can update the update target (reader and partition iterator).

Does this PR introduce any user-facing change?

No

How was this patch tested?

Existing tests.

Was this patch authored or co-authored using generative AI tooling?

No

dongjoon-hyun

+1, LGTM. Thank you, @viirya .
(Pending CIs).

dongjoon-hyun · 2025-07-15T20:44:31Z

cc @peter-toth

mridulm

LGTM

mridulm · 2025-07-16T02:03:03Z

+ new PartitionMetricCallback[ColumnarBatch](customMetrics)
+ } else {
+ new PartitionMetricCallback[InternalRow](customMetrics)
+ }


nit: We dont really need the if/else here ... the type param gets ignored.

Removed. Thanks.

…partitions in task completion listeners for metric update ### What changes were proposed in this pull request? This patch adds only one task completion listener for metric updating for `DataSourceRDD`, instead of adding separate one for each partition iterator. ### Why are the changes needed? For each partition iterator, currently we add one task completion listener used to update final metrics if the task is stopped early. In the listener, the reader and iterator are held. So if the partition is normally exhausted, the references cannot be released early. It is a problem especially if the references are heavy as reported by apache/iceberg#13297. Since the purpose of the callback is to update the final metrics if the task is stopped early, it means that we only need to do it for the last partition iterator. So we don't need set up a listener for each partition iterator. Thus, we can set up just one listener for all iterator. Once we advance to next partition iterator, we can update the update target (reader and partition iterator). ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closes#51503 from viirya/fix_metric_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit cea0051) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…partitions in task completion listeners for metric update This patch adds only one task completion listener for metric updating for `DataSourceRDD`, instead of adding separate one for each partition iterator. For each partition iterator, currently we add one task completion listener used to update final metrics if the task is stopped early. In the listener, the reader and iterator are held. So if the partition is normally exhausted, the references cannot be released early. It is a problem especially if the references are heavy as reported by apache/iceberg#13297. Since the purpose of the callback is to update the final metrics if the task is stopped early, it means that we only need to do it for the last partition iterator. So we don't need set up a listener for each partition iterator. Thus, we can set up just one listener for all iterator. Once we advance to next partition iterator, we can update the update target (reader and partition iterator). No Existing tests. No Closes#51503 from viirya/fix_metric_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit cea0051) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

viirya · 2025-07-16T07:22:18Z

Merged to master/branch-4.0/branch-3.5.

Thank you @dongjoon-hyun @szehon-ho @mridulm

…cing to next file in ParquetPartitionReaderFactory ### What changes were proposed in this pull request? This patch adds only one task completion listener for closing iterators in `ParquetPartitionReaderFactory`, instead of adding separate one for each file iterator. ### Why are the changes needed? For each file iterator, currently we add one task completion listener used to update closing the iterator when the task is finished. In the listener, the iterator reference is held. So if the file is normally exhausted, the reference cannot be released early. It is a problem especially if the reference are heavy as reported by apache/iceberg#13297. Similar to #51503, we don't need set up a listener for each file iterator. Thus, we can set up just one listener for all iterator. Once we advance to next file, we can update the update target to new iterator. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closes#51528 from viirya/fix_iter_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…cing to next file in ParquetPartitionReaderFactory ### What changes were proposed in this pull request? This patch adds only one task completion listener for closing iterators in `ParquetPartitionReaderFactory`, instead of adding separate one for each file iterator. ### Why are the changes needed? For each file iterator, currently we add one task completion listener used to update closing the iterator when the task is finished. In the listener, the iterator reference is held. So if the file is normally exhausted, the reference cannot be released early. It is a problem especially if the reference are heavy as reported by apache/iceberg#13297. Similar to #51503, we don't need set up a listener for each file iterator. Thus, we can set up just one listener for all iterator. Once we advance to next file, we can update the update target to new iterator. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closes#51528 from viirya/fix_iter_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 197c9d6) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…partitions in task completion listeners for metric update ### What changes were proposed in this pull request? This patch adds only one task completion listener for metric updating for `DataSourceRDD`, instead of adding separate one for each partition iterator. ### Why are the changes needed? For each partition iterator, currently we add one task completion listener used to update final metrics if the task is stopped early. In the listener, the reader and iterator are held. So if the partition is normally exhausted, the references cannot be released early. It is a problem especially if the references are heavy as reported by apache/iceberg#13297. Since the purpose of the callback is to update the final metrics if the task is stopped early, it means that we only need to do it for the last partition iterator. So we don't need set up a listener for each partition iterator. Thus, we can set up just one listener for all iterator. Once we advance to next partition iterator, we can update the update target (reader and partition iterator). ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closesapache#51503 from viirya/fix_metric_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 007d593) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…cing to next file in ParquetPartitionReaderFactory ### What changes were proposed in this pull request? This patch adds only one task completion listener for closing iterators in `ParquetPartitionReaderFactory`, instead of adding separate one for each file iterator. ### Why are the changes needed? For each file iterator, currently we add one task completion listener used to update closing the iterator when the task is finished. In the listener, the iterator reference is held. So if the file is normally exhausted, the reference cannot be released early. It is a problem especially if the reference are heavy as reported by apache/iceberg#13297. Similar to apache#51503, we don't need set up a listener for each file iterator. Thus, we can set up just one listener for all iterator. Once we advance to next file, we can update the update target to new iterator. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closesapache#51528 from viirya/fix_iter_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit dcdb482) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…cing to next file in ParquetPartitionReaderFactory ### What changes were proposed in this pull request? This patch adds only one task completion listener for closing iterators in `ParquetPartitionReaderFactory`, instead of adding separate one for each file iterator. ### Why are the changes needed? For each file iterator, currently we add one task completion listener used to update closing the iterator when the task is finished. In the listener, the iterator reference is held. So if the file is normally exhausted, the reference cannot be released early. It is a problem especially if the reference are heavy as reported by apache/iceberg#13297. Similar to apache#51503, we don't need set up a listener for each file iterator. Thus, we can set up just one listener for all iterator. Once we advance to next file, we can update the update target to new iterator. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No Closesapache#51528 from viirya/fix_iter_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 197c9d6) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

…partitions in task completion listeners for metric update This patch adds only one task completion listener for metric updating for `DataSourceRDD`, instead of adding separate one for each partition iterator. For each partition iterator, currently we add one task completion listener used to update final metrics if the task is stopped early. In the listener, the reader and iterator are held. So if the partition is normally exhausted, the references cannot be released early. It is a problem especially if the references are heavy as reported by apache/iceberg#13297. Since the purpose of the callback is to update the final metrics if the task is stopped early, it means that we only need to do it for the last partition iterator. So we don't need set up a listener for each partition iterator. Thus, we can set up just one listener for all iterator. Once we advance to next partition iterator, we can update the update target (reader and partition iterator). No Existing tests. No Closesapache#51503 from viirya/fix_metric_callback. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit cea0051) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

github-actionsBot added the SQL label Jul 15, 2025

viirya force-pushed the fix_metric_callback branch from 402e28b to b03c749Compare July 15, 2025 20:14

viirya mentioned this pull request Jul 15, 2025
Spark: Doing a Coalesce and foreachpartitions in spark directly on an iceberg table is leaking memory heavy iterators apache/iceberg#13297
Closed
3 tasks

viirya changed the title ~~[SPARK-XXXXX][SQL] Add only one task completion listener for metric update callback~~[SPARK-XXXXX][SQL] Don't hold reader and iterator references for all partitions in task completion listenersJul 15, 2025

viirya changed the title ~~[SPARK-XXXXX][SQL] Don't hold reader and iterator references for all partitions in task completion listeners~~[SPARK-XXXXX][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric updateJul 15, 2025

viirya changed the title ~~[SPARK-XXXXX][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric update~~[SPARK-52809][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric updateJul 15, 2025

dongjoon-hyun approved these changes Jul 15, 2025
View reviewed changes

szehon-ho approved these changes Jul 15, 2025
View reviewed changes

Add only one task completion listener for metric update callback
f9645ad

viirya force-pushed the fix_metric_callback branch from b03c749 to f9645adCompare July 15, 2025 21:46

dongjoon-hyun approved these changes Jul 15, 2025
View reviewed changes

mridulm approved these changes Jul 16, 2025
View reviewed changes

mridulm reviewed Jul 16, 2025
View reviewed changes

Remove type param
47fdc36

viirya closed this in cea0051Jul 16, 2025

viirya deleted the fix_metric_callback branch July 16, 2025 07:22

viirya mentioned this pull request Jul 17, 2025
[SPARK-52516][SQL] Don't hold previous iterator reference after advancing to next file in ParquetPartitionReaderFactory #51528
Closed

jkolash mentioned this pull request Jul 17, 2025
This hacky fix shows an improvement in memory usage for iceberg#13297 apache/iceberg#13298
Closed

szehon-ho mentioned this pull request Feb 19, 2026
[SPARK-55535][SPARK-55092][SQL] Refactor KeyGroupedPartitioning and Storage Partition Join #54330
Closed

peter-toth mentioned this pull request Feb 20, 2026
[SPARK-55619][SQL] Fix custom metrics in case of coalesced partitions #54396
Closed

Navigation Menu

Uh oh!

[SPARK-52809][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric update - #51503

[SPARK-52809][SQL] Don't hold reader and iterator references for all partitions in task completion listeners for metric update#51503
viirya wants to merge 2 commits into
apache:masterfrom
viirya:fix_metric_callback

viirya commented Jul 15, 2025•
edited
Loading

Uh oh!

dongjoon-hyun left a comment

Uh oh!

dongjoon-hyun commented Jul 15, 2025

Uh oh!

mridulm left a comment

Uh oh!

mridulmJul 16, 2025•
edited
Loading

Uh oh!

viiryaJul 16, 2025

Uh oh!

viirya commented Jul 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

viirya commented Jul 15, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun commented Jul 15, 2025

Uh oh!

mridulm left a comment

Choose a reason for hiding this comment

Uh oh!

mridulmJul 16, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viiryaJul 16, 2025

Choose a reason for hiding this comment

Uh oh!

viirya commented Jul 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

viirya commented Jul 15, 2025•
edited
Loading

mridulmJul 16, 2025•
edited
Loading