[jackboyla]

Related to #8793.

IterableDataset.iter(...).close() can leave nested generators alive while the caller keeps the closed iterator. On Python 3.12 this leaves shuffled streaming sources and their files open until that iterator is deleted. Return an iterator that closes and releases its owned batch generator, so early close releases the source even when it remains in scope.

Parquet scanning keeps its existing threading, prebuffering and cross-row-group reads.

Validation:

• A real-file regression fails on current main for both Python and Arrow batch formats, then passes with this change. It also checks repeated close and iteration after close.

• Python 3.12.3 / PyArrow 25.0.1: 727 passed, 77 skipped, 5 deselected across iterable datasets, Parquet, Parquet I/O, streaming download and distributed tests. Skips require optional frameworks or slow tests; deselections cover Torch interleave cases and remote GeoParquet.

• Two FineWeb-HQ files, text projection, shuffle, 16 examples, early close with the iterator retained: clean shutdown on released PyArrow 25. Five controlled public-iterator shutdown runs with delayed reads also exit cleanly.

• Reconstructed larger Hub workload with both fixes: interleaved FineWeb-HQ and FineWeb2-HQ (French), ten files per source, shuffle buffer 8,192 and batches of 4,096. Two runs each processed 344,064 examples / 1,508,773,700 UTF-8 text bytes and exited cleanly with empty stderr. All output hashes matched the baseline.

• The matching unmodified-Datasets + CPU-only Arrow PR baseline aborted with the exact fatal PyGILState_Release error after a partial-read connection error during shutdown. Combined-fix close took 15.6–16.1 seconds, with peak RSS around 3 GB. These runs use full local Arrow source builds and reconstruct the reported scale; they are not the reporter's exact script or a released-PyArrow-25 large-scale comparison.

• Controlled public HTTP scan, three samples: 0.674 → 0.689 seconds median, 5 GETs in both versions, identical rows, batches and table bytes. Ruff lint/format and diff whitespace checks pass.

Draft: awaiting the unmerged upstream Arrow changes. A separate Arrow shutdown failure remains: a failed range read can leave other prebuffer reads active after reader close. The Arrow extension to apache/arrow#54467 waits for CPU and I/O pools together; full local source builds pass that regression and the combined Datasets shutdown checks. This ownership fix alone does not resolve every shutdown crash in #8793.

Compatibility: the return value supports iteration and close(); it no longer exposes generator-only send, throw or gi_* attributes. Those are not part of the documented API.

[lhoestq]

use_threads=False can cause important slow downs when it comes to streaming parquet datasets, I would look for a better fix. I'm pretty sure you will see a difference if you run a benchmark

[jackboyla]

thanks for that 🙇 converting back to draft until i know more here

I compared upstream with this branch using a 64 MiB file with 64 row groups over an HTTP range server with 25 ms per GET. With fsspec's default HTTP cache, three-sample medians were 0.165 → 0.730 seconds, with 3 → 25 GETs. A variant that restores threaded decoding but closes a reader for each row group still took 0.727 seconds / 25 GETs, so the lost coalescing across row groups also matters. A bounded FineWeb-HQ text-only scan showed little difference (27.6 → 28.4 seconds for 50,000 rows); that doesn't rule out the regression in the controlled case.

Further checks also found a remaining correctness gap: a failed range read can leave other prebuffer reads running even on this branch. I reduced that to a standalone PyArrow 25.0.1 reproducer; reader.close() returns while a read is active, followed by the same fatal shutdown error.

I don't think this revision should merge as it stands. The next direction is to fix pending prebuffer cleanup in Arrow while keeping threaded scans and cross-group coalescing.

[jackboyla]

Thanks for flagging the slowdown! I replaced the synchronous workaround with an iterator cleanup fix that keeps threading and cross-row-group reads intact. The benchmark now shows the same request count and similar timing.

There’s still a separate Arrow shutdown bug; I’ve proposed an extension to apache/arrow#54467. With both fixes, two matching 1.5 GB streaming runs exited cleanly, while the baseline crashed.

Keeping this in draft until the Arrow changes land.

[lhoestq]

great thx ! looking forward for the fix to be merged on the arrow side