Chore: Do not return empty record batches from streams #13794

mertak-synnada · 2024-12-16T14:44:58Z

Which issue does this PR close?

Closes #.

Rationale for this change

Some plans generate and emit empty RecordBatches. Since this behavior is unnecessary and incurs costs during stream execution, this change implements a solution for handling them in the streaming process.

The main intention is, to avoid returning empty RecordBatches, acting them as technical debt, and replacing them with Option<RecordBatch> signatures, to reflect the behavior in design.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

…e possible

Dandandan · 2024-12-16T15:38:10Z

datafusion/physical-plan/src/windows/window_agg_exec.rs

-                None => self.compute_aggregates(),
-            };
+                None => {
+                    let batch = concat_batches(&self.input.schema(), &self.batches)?;


This change moves concat_batches out of elapsed_compute

I think we should fix this prior to merge

alamb

Thank you @mertak-synnada -- this looks like a good change to me

I left some small style suggestions. My only real concern with this PR is that there are no tests and thus someone may change this behavior in the future and we likely wouldn't catch it in code review.

Thus, I would recommend creating some sort of unit test that would fail if the AggregateExec is changed so that it generates empty RecordBatches again.

alamb · 2024-12-16T19:19:57Z

datafusion/physical-plan/src/aggregates/row_hash.rs

+                                let Some(batch) = extract_ok!(self.emit(to_emit, false))
+                                else {
+                                    break 'reading_input;
+                                };
+                                self.exec_state = ExecutionState::ProducingOutput(batch);


Nit: I think I would find this easier to read if it it avoid a redundant break:

Suggested change

let Some(batch) = extract_ok!(self.emit(to_emit, false))

else {

break 'reading_input;

};

self.exec_state = ExecutionState::ProducingOutput(batch);

if let let Some(batch) = extract_ok!(self.emit(to_emit, false)) {

self.exec_state = ExecutionState::ProducingOutput(batch);

}

alamb · 2024-12-16T19:20:13Z

datafusion/physical-plan/src/aggregates/row_hash.rs

                                timer.done();
+                                let Some(batch) = extract_ok!(self.emit(to_emit, false))


same comment as above

alamb · 2024-12-16T19:25:44Z

datafusion/physical-plan/src/aggregates/row_hash.rs

        let schema = if spilling {
            Arc::clone(&self.spill_state.spill_schema)
        } else {
            self.schema()
        };
        if self.group_values.is_empty() {
-            return Ok(RecordBatch::new_empty(schema));
+            return Ok(None);


alamb · 2024-12-16T19:26:15Z

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -937,7 +944,8 @@ impl GroupedHashAggregateStream {
        // over the target memory size after emission, we can emit again rather than returning Err.
        let _ = self.update_memory_reservation();
        let batch = RecordBatch::try_new(schema, output)?;
-        Ok(batch)
+        debug_assert!(batch.num_rows() > 0);


Maybe it would be good to document in comments somewhere the expectation / behavior that no empty record batches are produced

take concat_batches into compute_aggregates function again

mertak-synnada · 2024-12-17T08:35:12Z

Thanks for the review, @alamb and @Dandandan! I tried to implement your suggestions, including adding some unit tests and assertions for some existing tests. I hope it's in a better state now.

ozankabak

LGTM in general, left some minor feedback

ozankabak · 2024-12-17T14:00:01Z

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -768,6 +776,9 @@ impl Stream for GroupedHashAggregateStream {
                        let output = batch.slice(0, size);
                        (ExecutionState::ProducingOutput(remaining), output)
                    };
+                    // Empty record batches should not be emitted.
+                    // They need to be treated as  [`Option<RecordBatch>`]es and handle separately


Suggested change

// They need to be treated as [`Option<RecordBatch>`]es and handle separately

// They need to be treated as [`Option<RecordBatch>`]es and handled separately

ozankabak · 2024-12-17T14:00:52Z

datafusion/physical-plan/src/sorts/partial_sort.rs

@@ -407,6 +407,9 @@ impl PartialSortStream {
                self.is_closed = true;
            }
        }
+        // Empty record batches should not be emitted.
+        // They need to be treated as [`Option<RecordBatch>`]es and handle separately


Suggested change

// They need to be treated as [`Option<RecordBatch>`]es and handle separately

// They need to be treated as [`Option<RecordBatch>`]es and handled separately

ozankabak · 2024-12-17T14:01:24Z

datafusion/physical-plan/src/unnest.rs

+                    self.metrics.output_rows.add(result_batch.num_rows());
+
+                    // Empty record batches should not be emitted.
+                    // They need to be treated as  [`Option<RecordBatch>`]es and handle separately


Suggested change

// They need to be treated as [`Option<RecordBatch>`]es and handle separately

// They need to be treated as [`Option<RecordBatch>`]es and handled separately

ozankabak · 2024-12-17T14:04:20Z

datafusion/physical-plan/src/unnest.rs

+                let temp_result = list_unnest_at_level(
                    input,
                    list_type_columns,
                    &mut temp_unnested_result,
                    depth,
                    options,
                )?;
-                if num_rows == 0 {
-                    return Ok(RecordBatch::new_empty(Arc::clone(schema)));
-                }
+                let Some(temp_result) = temp_result else {
+                    return Ok(None);
+                };


We can use the more succinct

let Some(temp_result) = list_unnest_at_level( input, list_type_columns, &mut temp_unnested_result, depth, options, )? else { continue; };

ozankabak · 2024-12-17T14:05:20Z

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

+                let batch = self.compute_aggregates()?;
+                if let Some(batch) = batch {
+                    return Poll::Ready(Some(Ok(batch)));
+                }


Can we say

Suggested change

let batch = self.compute_aggregates()?;

if let Some(batch) = batch {

return Poll::Ready(Some(Ok(batch)));

}

if let Some(batch) = self.compute_aggregates()? {

return Poll::Ready(Some(Ok(batch)));

}

ozankabak · 2024-12-17T14:05:51Z

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

+                let batch = self.compute_aggregates()?;
+                if let Some(batch) = batch {
+                    return Poll::Ready(Some(Ok(batch)));
+                }


Can we say

Suggested change

let batch = self.compute_aggregates()?;

if let Some(batch) = batch {

return Poll::Ready(Some(Ok(batch)));

}

if let Some(batch) = self.compute_aggregates()? {

return Poll::Ready(Some(Ok(batch)));

}

ozankabak · 2024-12-17T14:07:02Z

datafusion/physical-plan/src/windows/window_agg_exec.rs

+            let Some(result) = maybe_batch else {
+                return Poll::Ready(None);
            };


Can this go inside the None arm of the match above?

ozankabak · 2024-12-17T14:07:11Z

datafusion/physical-plan/src/windows/window_agg_exec.rs

-
-            return Poll::Ready(Some(result));
+            // Empty record batches should not be emitted.
+            // They need to be treated as  [`Option<RecordBatch>`]es and handle separately


Suggested change

// They need to be treated as [`Option<RecordBatch>`]es and handle separately

// They need to be treated as [`Option<RecordBatch>`]es and handled separately

ozankabak · 2024-12-18T07:42:17Z

I will go ahead and merge this in the next few hours unless new feedback arrives, since the reviews has been addressed and it is a fairly straightforward change

alamb · 2024-12-18T23:38:32Z

Thank you everyone!

mertak-synnada added 7 commits December 13, 2024 15:52

do not emit empty record batches in plans

e9cc59a

Merge branch 'refs/heads/apache_main' into chore/remove-empty-batches

30a6b97

change function signatures to Option<RecordBatch> if empty batches ar…

69fbe64

…e possible

format code

2c7b70d

shorten code

4b56e65

change list_unnest_at_level for returning Option value

fd44a6e

Merge branch 'refs/heads/apache_main' into chore/remove-empty-batches

4706bc3

github-actions bot added the physical-expr Physical Expressions label Dec 16, 2024

Dandandan reviewed Dec 16, 2024

View reviewed changes

alamb changed the title ~~Chore: Remove empty batches~~ Chore: do not return empty record batches from Aggregate operators Dec 16, 2024

alamb approved these changes Dec 16, 2024

View reviewed changes

mertak-synnada added 8 commits December 17, 2024 10:03

add documentation

72bdab9

take concat_batches into compute_aggregates function again

create unit test for row_hash.rs

1f9349f

add test for unnest

cd505f4

add test for unnest

0dccc81

add test for partial sort

466d697

add test for bounded window agg

29e1204

add test for window agg

f9c2bbe

Merge branch 'refs/heads/apache_main' into chore/remove-empty-batches

9978e4d

github-actions bot added the core Core DataFusion crate label Dec 17, 2024

mertak-synnada requested review from alamb and Dandandan December 17, 2024 10:19

ozankabak approved these changes Dec 17, 2024

View reviewed changes

mertak-synnada added 2 commits December 17, 2024 17:32

apply simplifications and fix typo

c59afd4

apply simplifications and fix typo

2730f81

mertak-synnada changed the title ~~Chore: do not return empty record batches from Aggregate operators~~ Chore: Do not return empty record batches from streams Dec 18, 2024

ozankabak merged commit 63ce486 into apache:main Dec 18, 2024
26 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Chore: Do not return empty record batches from streams #13794

Chore: Do not return empty record batches from streams #13794

mertak-synnada commented Dec 16, 2024 •

edited

Loading

Dandandan Dec 16, 2024

alamb Dec 16, 2024

alamb left a comment

alamb Dec 16, 2024

alamb Dec 16, 2024

alamb Dec 16, 2024

alamb Dec 16, 2024

mertak-synnada commented Dec 17, 2024

ozankabak left a comment

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak Dec 17, 2024

ozankabak commented Dec 18, 2024

alamb commented Dec 18, 2024

		timer.done();
		let Some(batch) = extract_ok!(self.emit(to_emit, false))

	// They need to be treated as [`Option<RecordBatch>`]es and handle separately
	// They need to be treated as [`Option<RecordBatch>`]es and handled separately

Chore: Do not return empty record batches from streams #13794

Chore: Do not return empty record batches from streams #13794

Conversation

mertak-synnada commented Dec 16, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mertak-synnada commented Dec 17, 2024

ozankabak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ozankabak commented Dec 18, 2024

alamb commented Dec 18, 2024

mertak-synnada commented Dec 16, 2024 •

edited

Loading