You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2023/07/19 10:38:35 UTC
[GitHub] [arrow-datafusion] alamb opened a new pull request, #7025: Minor: Refactor group by state machine
alamb opened a new pull request, #7025:
URL: https://github.com/apache/arrow-datafusion/pull/7025
# Which issue does this PR close?
Follow on to https://github.com/apache/arrow-datafusion/pull/6932
# Rationale for this change
I was a little unhappy about the `extract_ok!` macro and @tustvold
offered a suggestion on how to remove it: https://github.com/apache/arrow-datafusion/pull/6932/files#r1267309181
This PR is what I came up with . I am not really sure it is better but wanted to put it up in case others had an opinion
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
Pull out the state machine transition logic in group by into some new functions and remove the `extract_ok!` macro
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
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@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #7025: Minor: Refactor group by state machine
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #7025:
URL: https://github.com/apache/arrow-datafusion/pull/7025#discussion_r1268175366
##########
datafusion/core/src/physical_plan/aggregates/row_hash.rs:
##########
@@ -342,69 +332,92 @@ impl Stream for GroupedHashAggregateStream {
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
loop {
- let exec_state = self.exec_state.clone();
- match exec_state {
+ let res = match self.exec_state.clone() {
ExecutionState::ReadingInput => {
- match ready!(self.input.poll_next_unpin(cx)) {
- // new batch to aggregate
- Some(Ok(batch)) => {
- let timer = elapsed_compute.timer();
- // Do the grouping
- extract_ok!(self.group_aggregate_batch(batch));
-
- // If we can begin emitting rows, do so,
- // otherwise keep consuming input
- assert!(!self.input_done);
- let to_emit = self.group_ordering.emit_to();
-
- if let Some(to_emit) = to_emit {
- let batch =
- extract_ok!(self.create_batch_from_map(to_emit));
- self.exec_state = ExecutionState::ProducingOutput(batch);
- }
- timer.done();
- }
- Some(Err(e)) => {
- // inner had error, return to caller
- return Poll::Ready(Some(Err(e)));
- }
- None => {
- // inner is done, emit all rows and switch to producing output
- self.input_done = true;
- self.group_ordering.input_done();
- let timer = elapsed_compute.timer();
- let batch =
- extract_ok!(self.create_batch_from_map(EmitTo::All));
- self.exec_state = ExecutionState::ProducingOutput(batch);
- timer.done();
- }
- }
+ let input = ready!(self.input.poll_next_unpin(cx));
+ self.next_input(input, &elapsed_compute)
}
ExecutionState::ProducingOutput(batch) => {
- // slice off a part of the batch, if needed
- let output_batch = if batch.num_rows() <= self.batch_size {
- if self.input_done {
- self.exec_state = ExecutionState::Done;
- } else {
- self.exec_state = ExecutionState::ReadingInput
- }
- batch
- } else {
- // output first batch_size rows
- let num_remaining = batch.num_rows() - self.batch_size;
- let remaining = batch.slice(self.batch_size, num_remaining);
- self.exec_state = ExecutionState::ProducingOutput(remaining);
- batch.slice(0, self.batch_size)
- };
- return Poll::Ready(Some(Ok(
- output_batch.record_output(&self.baseline_metrics)
- )));
+ let output_batch = self.next_output(batch);
+ return Poll::Ready(Some(Ok(output_batch)));
}
ExecutionState::Done => return Poll::Ready(None),
+ };
+
+ // return error, otherwise loop again
+ if let Err(e) = res {
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
+ }
+}
+
+impl GroupedHashAggregateStream {
+ /// Processes the next batch of input, updating
+ /// self.execution_state appropriately
+ fn next_input(
+ &mut self,
+ input: Option<Result<RecordBatch>>,
+ elapsed_compute: &Time,
+ ) -> Result<()> {
+ assert!(matches!(&self.exec_state, ExecutionState::ReadingInput));
+ match input {
+ // new batch to aggregate
+ Some(batch) => {
+ let batch = batch?;
+ let timer = elapsed_compute.timer();
+ // Do the grouping
+ self.group_aggregate_batch(batch)?;
+
+ // If we can begin emitting rows, do so,
+ // otherwise keep consuming input
+ assert!(!self.input_done);
+ let to_emit = self.group_ordering.emit_to();
+
+ if let Some(to_emit) = to_emit {
+ let batch = self.create_batch_from_map(to_emit)?;
+ self.exec_state = ExecutionState::ProducingOutput(batch);
+ }
+ timer.done();
+ }
+ None => {
+ // inner is done, emit all rows and switch to producing output
+ self.input_done = true;
+ self.group_ordering.input_done();
+ let timer = elapsed_compute.timer();
+ let batch = self.create_batch_from_map(EmitTo::All)?;
+ self.exec_state = ExecutionState::ProducingOutput(batch);
+ timer.done();
}
}
+ Ok(())
+ }
+
+ /// Produces the next batch of output from ExecutionState::ProductingOuptut, updating self.execution_state
Review Comment:
```suggestion
/// Produces the next batch of output from ExecutionState::ProducingOuptut, updating self.execution_state
```
--
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@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #7025: Minor: Refactor group by state machine
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7025:
URL: https://github.com/apache/arrow-datafusion/pull/7025#discussion_r1267913203
##########
datafusion/core/src/physical_plan/aggregates/row_hash.rs:
##########
@@ -322,16 +322,6 @@ fn create_group_accumulator(
}
}
-/// Extracts a successful Ok(_) or returns Poll::Ready(Some(Err(e))) with errors
-macro_rules! extract_ok {
Review Comment:
This macro was removed, but I think breaking the main state loop change into two different functions makes is slightly harder to follow the logic.
Thus I am inclined not to pursue this PR anymore
--
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@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #7025: Minor: Refactor group by state machine
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #7025:
URL: https://github.com/apache/arrow-datafusion/pull/7025#issuecomment-1646622838
I think I am inclined to leave this code alone for now as it has already undergone a bunch of churn and this is just stylistic and I want to move on to other things. If anyone else feels differently, please feel free to open a new PR
--
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@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb closed pull request #7025: Minor: Refactor group by state machine
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb closed pull request #7025: Minor: Refactor group by state machine
URL: https://github.com/apache/arrow-datafusion/pull/7025
--
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@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org