You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "feniljain (via GitHub)" <gi...@apache.org> on 2025/12/15 19:29:41 UTC
[PR] wip: integrate batch coalescer with async fn exec [datafusion]
feniljain opened a new pull request, #19342:
URL: https://github.com/apache/datafusion/pull/19342
## Which issue does this PR close?
- Closes #19331
## What changes are included in this PR?
Currently haven't included changes to remove it from optimizer, will do it in a subsequent PR as I see everyone else has done the same?
Mostly took code from other PRs and just stitched it together to work with AsyncFnExec
## Are these changes tested?
In Progress
## Are there any user-facing changes?
No
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2635240485
##########
datafusion/physical-plan/src/async_func.rs:
##########
@@ -221,6 +234,49 @@ impl ExecutionPlan for AsyncFuncExec {
}
}
+struct CoalesceInputStream {
+ input_stream: Pin<Box<dyn RecordBatchStream + Send>>,
+ batch_coalescer: LimitedBatchCoalescer,
+}
+
+impl Stream for CoalesceInputStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let mut completed = false;
+
Review Comment:
I had omitted it as we were not doing it already, but if it makes sense then sure! Pushing a commit for it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] wip: integrate batch coalescer with async fn exec [datafusion]
Posted by "jizezhang (via GitHub)" <gi...@apache.org>.
jizezhang commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2635203037
##########
datafusion/physical-plan/src/async_func.rs:
##########
@@ -221,6 +234,49 @@ impl ExecutionPlan for AsyncFuncExec {
}
}
+struct CoalesceInputStream {
+ input_stream: Pin<Box<dyn RecordBatchStream + Send>>,
+ batch_coalescer: LimitedBatchCoalescer,
+}
+
+impl Stream for CoalesceInputStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let mut completed = false;
+
Review Comment:
Do we want to track `elapsed_compute` time for coalescing? e.g. https://github.com/apache/datafusion/blob/887aa9f8c641cfcac4053da3dd5ad3c4b7ffe425/datafusion/physical-plan/src/coalesce_batches.rs#L296
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#issuecomment-3700106859
Thanks for patiently reviewing and helping me! ππ»
>BTW I think we'll need one more PR to remove the rule from CoalesceBatches (and maybe we can even remove the entire rule now...)
I will raise a PR for both of them after this gets merged :)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#issuecomment-3703714636
Thanks @feniljain and @jizezhang
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #19342:
URL: https://github.com/apache/datafusion/pull/19342
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#issuecomment-3703721653
> Thank you for patiently reviewing and helping me! ππ»
>
> > BTW I think we'll need one more PR to remove the rule from CoalesceBatches (and maybe we can even remove the entire rule now...)
>
> I will raise a PR for both of them after this gets merged. And also another PR to implement metrics in `AsyncFuncExec` :)
Thank you @feniljain . I filed a ticket to track this:
- https://github.com/apache/datafusion/issues/19591
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "jizezhang (via GitHub)" <gi...@apache.org>.
jizezhang commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2637249921
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
https://github.com/apache/datafusion/pull/19002 contains test, e.g.
https://github.com/apache/datafusion/blob/d8e68a404aa7955b4913bab1c658a4ec5c82f0fc/datafusion/physical-plan/src/repartition/mod.rs#L1753
Have you by any chance tried leveraging `TestMemoryExec` https://github.com/apache/datafusion/blob/d8e68a404aa7955b4913bab1c658a4ec5c82f0fc/datafusion/physical-plan/src/test.rs#L62 to create an input plan/stream to this exec and check output batch sizes? or it did not work?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] wip: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#issuecomment-3674432511
Hey @jizezhang ππ»
I think we can review it now :)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2635323310
##########
datafusion/physical-plan/src/async_func.rs:
##########
@@ -221,6 +234,49 @@ impl ExecutionPlan for AsyncFuncExec {
}
}
+struct CoalesceInputStream {
+ input_stream: Pin<Box<dyn RecordBatchStream + Send>>,
+ batch_coalescer: LimitedBatchCoalescer,
+}
+
+impl Stream for CoalesceInputStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let mut completed = false;
+
Review Comment:
It seems this was left as a `TODO` in general [here](https://github.com/apache/datafusion/blob/91cfb6990319a672195f97d0b2b179fa9634fec9/datafusion/physical-plan/src/async_func.rs#L181)
As all metrics need to be implemented, I think it will be better to do it in a separate PR? What do you think?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2637134301
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
@alamb I was playing around with this and have a small doubt:
Lets consider the case where input batches are smaller than coalescer's target batch size. Above code is implemented as a `then` on `input_stream`. Every iteration we need to return at least one batch. That would mean we will have to poll `input_stream` again in `then` itself. But as `input_stream` is already used we can not use it again in `then` block, this is causing an error and I don't see a good way out of it except wrapping it around a stream. Should we leave it as it is? Do help me if I am missing some other obvious thing here π
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2637134301
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
@alamb I was playing around with this and have a small doubt:
Lets consider the case where input batches are smaller than coalescer's target batch size. Above code is implemented as a `then` on `input_stream`. Every iteration we need to return at least one batch. That would mean we will have to poll `input_stream` again in `then` itself. But as `input_stream` is already used we can not use it again in `then` block, this is causing an error and I don't see a good way out of it except wrapping it around a stream. Should we leave it as it is? Do correct me if I am missing some other obvious thing here π
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2637243357
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
Regarding tests, I see we don't have tests where I have control over stream, which is required to calculate how many times did it get called. And I am not sure how hard would it be create a test with these conditions.
I tested this flow earlier by keeping a `DEBUG` log in stream flow of `AsyncFuncExec`, manipulating `target_batch_size` in [this](https://github.com/apache/datafusion/blob/d8e68a404aa7955b4913bab1c658a4ec5c82f0fc/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs#L33) test, this helped me confirm everything was working fine.
Also I checked other two PRs and none contained tests:
- https://github.com/apache/datafusion/pull/19002/
- https://github.com/apache/datafusion/pull/18604
I suspect tests were not added in them for the same reason I mentioned above, do correct me if I am missing a point, thanks! ππ»
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2635323310
##########
datafusion/physical-plan/src/async_func.rs:
##########
@@ -221,6 +234,49 @@ impl ExecutionPlan for AsyncFuncExec {
}
}
+struct CoalesceInputStream {
+ input_stream: Pin<Box<dyn RecordBatchStream + Send>>,
+ batch_coalescer: LimitedBatchCoalescer,
+}
+
+impl Stream for CoalesceInputStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let mut completed = false;
+
Review Comment:
It seems this was left as a `TODO` in general [here](https://github.com/apache/datafusion/blob/91cfb6990319a672195f97d0b2b179fa9634fec9/datafusion/physical-plan/src/async_func.rs#L181)
As all metrics need to be implemented, I think it will be better to do it in a separate PR? What do you think?
I will take up metrics PR after this.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#issuecomment-3699337499
> * Closes [Integrate `BatchCoalescer` into `AsyncFuncExec` and remove from `CoalesceBatches` optimization ruleΒ #19331](https://github.com/apache/datafusion/issues/19331)
BTW I think we'll need one more PR to remove the rule from `CoalesceBatches` (and maybe we can even remove the entire rule now...)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "feniljain (via GitHub)" <gi...@apache.org>.
feniljain commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2652831217
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
Thanks for the pointer! ππ»
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: integrate batch coalescer with async fn exec [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #19342:
URL: https://github.com/apache/datafusion/pull/19342#discussion_r2637066839
##########
datafusion/physical-plan/src/async_func.rs:
##########
Review Comment:
I think you could move the batch coalescer into this logic directly which would
1. Make this PR smaller
2. Be more efficient (we wouldn't end up with the overhead of another stream at runtime)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org