Skip to content

Commit 3a29d6b

Browse files
pepijnverluvaton
andauthored
Add minimal genarator-like stream implementation (apache#23530)
## Which issue does this PR close? None, related to PR apache#23407 ## Rationale for this change Manually writing `Stream` implementations can be quite tedious since it requires implementing the state machine of the stream yourself. This often results in hard to read code. The same problem exists with manual `Future` implementations and `async` was added to the Rust language to mitigate exactly this problem. Unfortunately there's no language level support yet to help with writing streams/generators. There are quite a few projects that attempt to fill this gap: - Tokio's [async_stream](https://docs.rs/async-stream/latest/async_stream/) provides a proc macro that recognises a `yield` keyword. This is a good implementation, but proc macros don't play nice with code formatting, increase compile time, and in this particular case prevent decomposition into smaller functions. - [async_fn_stream](https://docs.rs/async-fn-stream/latest/async_fn_stream/) provides an implementation of the same concept, but without the proc macro. Unfortunately this implementation chooses a heavier mechanism to communicate generated values back to the stream via a `SmallVec`. - [genawaiter](https://docs.rs/genawaiter/latest/genawaiter/sync/index.html) is a more general purpose generator library, but this can also be used to implement `Stream`s. This project does miss some of the ergonomics provided by the other two in the form of `try_` variants that help in implementing fallible streams. Since none of these variants seems like the ideal candidate, the best option might be to have a custom implementation of the concept tailored to the needs of the DataFusion project. This PR provides an initial draft of exactly that. ## What changes are included in this PR? - Adds `async_stream` and `async_try_stream` functions that create Stream implementations based on an async generator function. The initial implementation was inspired by the macro expansion produced by `async_stream`. The code was then adapted further taking inspiration from the two other libraries. Specifically, the `Emitter` terminology was taken from `async_fn_stream` and the `Arc<Mutex<Option<T>>>` value transfer mechanism was taken from `genawaiter`. `async_stream` uses a very light weight thread-local storage based mechanism to handle value transfer, but this felt too risky to use when the emitter is exposed. It makes sense in the Tokio implementation since the proc macro can manage control flow in a stricter way. I wasn't entirely sure if taking some inspiration from has code licensing implications. I applied ASLv2 for now on the added code. ## Are these changes tested? Test code was mainly adapted from the tokio implementation. ## Are there any user-facing changes? Yes, new public function `async_stream` and `async_try_stream` --------- Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com>
1 parent 21986cf commit 3a29d6b

6 files changed

Lines changed: 736 additions & 1 deletion

File tree

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ pbjson = { version = "0.9.0" }
187187
pbjson-types = "0.9"
188188
percent-encoding = "2.3"
189189
pin-project = "1"
190+
pin-project-lite = "^0.2.7"
190191
# Should match arrow-flight's version of prost.
191192
prost = "0.14.1"
192193
rand = "0.9"

datafusion/execution/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ log = { workspace = true }
6565
object_store = { workspace = true, features = ["fs"] }
6666
parking_lot = { workspace = true }
6767
parquet = { workspace = true, optional = true }
68+
pin-project-lite = { workspace = true }
6869
rand = { workspace = true }
6970
tempfile = { workspace = true }
7071
tokio = { workspace = true }

0 commit comments

Comments
 (0)