rluvaton commented on code in PR #23702:
URL: https://github.com/apache/datafusion/pull/23702#discussion_r3653123989
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -212,79 +214,122 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
result
}
- fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
- async_try_stream(|mut emitter| async move {
- // This vector contains the indices of the partitions that have
not started emitting yet.
- let mut uninitiated_partitions =
- (0..self.streams.partitions()).collect::<Vec<_>>();
+ async fn flush_in_progress(
+ &mut self,
+ mut emitter: TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if self.in_progress.is_empty() {
+ return Ok(());
+ }
- poll_fn(|cx| self.initialize_all_partitions(&mut
uninitiated_partitions, cx))
- .await?;
+ let elapsed_compute = self.metrics.elapsed_compute().clone();
+ let mut timer = elapsed_compute.timer();
+
+ // When `build_record_batch()` hits an i32 offset overflow (e.g.
+ // combined string offsets exceed 2 GB), it emits a partial batch
+ // and keeps the remaining rows in `self.in_progress.indices`.
+ // Drain those leftover rows before terminating the stream,
+ // otherwise they would be silently dropped.
+ // Repeated overflows are fine — each poll emits another partial
+ // batch until `in_progress` is fully drained.
+ while let Some(batch) = self.emit_in_progress_batch()? {
+ drop(timer);
+ emitter.emit(batch).await;
+ timer = elapsed_compute.timer();
+ }
- assert_eq!(uninitiated_partitions.len(), 0);
+ Ok(())
+ }
- // If there are no more uninitiated partitions, set up the loser
tree and continue
- // to the next phase.
+ fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
+ async_try_stream(|mut emitter| async move {
+ assert!(
+ self.fetch.is_none_or(|fetch| fetch != 0),
Review Comment:
I wanted to have that on top to make it close to the code that depend on
that, but removing since redundant
--
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]
