Skip to content

fix: make Python record batch reader cancellable - #885

Merged
JingsongLi merged 1 commit into
apache:mainfrom
JingsongLi:codex/native-read-alignment-followup
Sep 20, 2026
Merged

JingsongLi merged 1 commit into
apache:mainfrom
JingsongLi:codex/native-read-alignment-followup

Conversation

@JingsongLi

Copy link
Copy Markdown
Contributor

Summary

  • add an idempotent close API to the Python RecordBatchReader
  • interrupt an in-flight async next operation without waiting on the stream mutex
  • release the native stream promptly and document the API in the Python type stub
  • cover pending-read cancellation plus Python close/EOF behavior

Motivation

Parallel PyPaimon reads may close after one reader produces a batch while another reader is still blocked in its first native read. The Python executor cannot finish until that in-flight Rust operation is cancellable. The close path uses an atomic flag and Tokio Notify, so it does not contend on the mutex held by next_record_batch.

Tests

  • cargo test -p pypaimon_rust record_batch_reader_close_interrupts_pending_next --locked
  • cargo clippy -p pypaimon_rust --all-targets --locked -- -D warnings
  • python -m pytest -q bindings/python/tests/test_read.py (82 passed)

@leaves12138 leaves12138 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

LGTM. Reviewed the stream mutex, atomic closed state, and stored Notify permit: the close path does not wait for an in-flight next, and the permit covers the gap between the closed check and notification polling.

Validation against this head:

  • 77 Rust binding tests passed, including three additional local tests for close-before-read, concurrent next/repeated close, and 200 close/first-poll races with stream-drop assertions.
  • Rebuilt the Python extension and ran test_read.py: 82 passed.
  • An additional end-to-end Python test held a native S3 request pending against a local HTTP endpoint; close interrupted read_next_batch before the endpoint was released, and repeated close/subsequent EOF behaved correctly.

No blocking correctness or lifecycle issues found in this change.

@JingsongLi
JingsongLi merged commit f33dee1 into apache:main Sep 20, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants