pepijnve commented on code in PR #23407:
URL: https://github.com/apache/datafusion/pull/23407#discussion_r3555728322
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -171,24 +152,40 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
in_progress: BatchBuilder::new(schema, stream_count, batch_size,
reservation),
streams,
metrics,
- done: false,
- drain_in_progress_on_done: false,
cursors: (0..stream_count).map(|_| None).collect(),
prev_cursors: (0..stream_count).map(|_| None).collect(),
round_robin_tie_breaker_mode: false,
num_of_polled_with_same_value: vec![0; stream_count],
current_reset_epoch: 0,
poll_reset_epochs: vec![0; stream_count],
loser_tree: vec![],
- loser_tree_adjusted: false,
batch_size,
fetch,
produced: 0,
- uninitiated_partitions: (0..stream_count).collect(),
enable_round_robin_tie_breaker,
}
}
+ pub(crate) fn into_stream(self) -> SendableRecordBatchStream
+ where
+ C: 'static,
+ {
+ let schema_clone = Arc::clone(self.in_progress.schema());
+
+ let cloned_metrics = self.metrics.clone();
+
+ let mut this = self;
+ let stream = Gen::new(|co| async move {
Review Comment:
A possible alternative to both could be a DIY solution for DataFusion. I
tried to sketch out an alternative loosely based on the various existing
implementations
https://github.com/pepijnve/datafusion/blob/4b4954cd50a26101d4e40d25c13d9068b00b8d2c/datafusion/execution/src/async_stream.rs
There's probably still room for improvement here, the main thing I wanted to
get across is that this looks much nicer imo at the call site in
SortPreservingMergeStream (see
https://github.com/pepijnve/datafusion/blob/4b4954cd50a26101d4e40d25c13d9068b00b8d2c/datafusion/physical-plan/src/sorts/merge.rs#L214)
The main downside I see is that without a proc macro it's really hard to
make this entirely abuse proof. I started with the thread local approach used
by the `async-stream` macro, but that's unsound when the emitter is exposed. I
went for the `Arc<Mutex>` then rather than the `SmallVec` that
`async-fn-stream` uses.
Doesn't have to be this implementation, but perhaps something like this
could be a "have your cake and eat it" solution?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
