feat: Support sliding window queries for MedianAccumulator by implementing `retract_batch` by petern48 · Pull Request #19278 · apache/datafusion

petern48 · 2025-12-11T07:20:44Z

Which issue does this PR close?

Closes Add retract_batch method for median accumulator #7664

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Added tests

Are there any user-facing changes?

Computing the median() window is now supported instead of throwing an error

datafusion/functions-aggregate/src/median.rs

2010YOUY01

This is great, thank you! Left some suggestions.

This PR adds a simple working solution, and it’s quite interesting to figure out how to retract efficiently for large windows 🤔

datafusion/functions-aggregate/src/median.rs

2010YOUY01 · 2025-12-11T10:49:59Z

datafusion/sqllogictest/test_files/aggregate.slt

+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING


I recommend to test different window frames like UNBOUNDED PRECEDING/FOLLOWING

I wasn't familiar with these before, but this was a great idea! It helped me find and understand a bug.

… for the bug in aggregate.slt

petern48 · 2025-12-13T20:31:09Z

datafusion/sqllogictest/test_files/aggregate.slt

+# median_non_sliding_window
+query ITRRRR
+SELECT
+    timestamp,
+    tags,
+    value,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
+    ) AS value_median_unbounded_preceding,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+    ) AS value_median_unbounded_both,


For UNBOUNDED FOLLOWING, an error is raised when retract_batch() isn't implemented. I found that queries with UNBOUNDED PRECEDING do not trigger this and instead return incorrect results. I assume this is a bug, right? If so, I can file a ticket.

For example, if you remove the UNBOUNDED FOLLOWING case right below my comment here, and try the query on main, I get this diff instead of an error.

Results Diff
``` [Diff] (-expected|+actual) 1 tag1 10 10 30 - 2 tag1 20 15 30 - 3 tag1 30 20 30 - 4 tag1 40 25 30 - 5 tag1 50 30 30 + 2 tag1 20 20 30 + 3 tag1 30 30 30 + 4 tag1 40 40 30 + 5 tag1 50 50 30 1 tag2 60 60 80 - 2 tag2 70 65 80 - 3 tag2 80 70 80 - 4 tag2 90 75 80 - 5 tag2 100 80 80 + 2 tag2 70 70 80 + 3 tag2 80 80 80 + 4 tag2 90 90 80 + 5 tag2 100 100 80 ```

There is this quite informative comment which seems to explain why this is the case:

datafusion/datafusion/physical-expr/src/aggregate.rs

Lines 490 to 538 in befaf93

// Accumulators that have window frame startings different

// than `UNBOUNDED PRECEDING`, such as `1 PRECEDING`, need to

// implement retract_batch method in order to run correctly

// currently in DataFusion.

//

// If this `retract_batches` is not present, there is no way

// to calculate result correctly. For example, the query

//

// ```sql

// SELECT

// SUM(a) OVER(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS sum_a

// FROM

// t

// ```

//

// 1. First sum value will be the sum of rows between `[0, 1)`,

//

// 2. Second sum value will be the sum of rows between `[0, 2)`

//

// 3. Third sum value will be the sum of rows between `[1, 3)`, etc.

//

// Since the accumulator keeps the running sum:

//

// 1. First sum we add to the state sum value between `[0, 1)`

//

// 2. Second sum we add to the state sum value between `[1, 2)`

// (`[0, 1)` is already in the state sum, hence running sum will

// cover `[0, 2)` range)

//

// 3. Third sum we add to the state sum value between `[2, 3)`

// (`[0, 2)` is already in the state sum). Also we need to

// retract values between `[0, 1)` by this way we can obtain sum

// between [1, 3) which is indeed the appropriate range.

//

// When we use `UNBOUNDED PRECEDING` in the query starting

// index will always be 0 for the desired range, and hence the

// `retract_batch` method will not be called. In this case

// having retract_batch is not a requirement.

//

// This approach is a a bit different than window function

// approach. In window function (when they use a window frame)

// they get all the desired range during evaluation.

if !accumulator.supports_retract_batch() {

return not_impl_err!(

"Aggregate can not be used as a sliding accumulator because \

`retract_batch` is not implemented: {}",

self.name

);

}

I think we should file a ticket, the previous impl should be able to handle unbounded preceding as @Jefffrey explained, and the inconsistent results is likely to indicate a bug.

Ah, so you're saying unbounded preceding is supposed to work even without retract_batch() implemented. I was originally under the impression that it wasn't, but no that makes total sense now.

In that case, I think this PR is already fixes the bug, so there's no need to submit an issue for that. I mentioned in this comment that passing mut instead of clearing state with take() (81ced74) fixes the results in the mod.rs test. I've verified this by copying that change (81ced74) over to main and testing it, and the results for that test change. It's completely unrelated to the new support for retract_batch(). We just have an integer overflow issue remaining, which I've submitted an issue for.

Yes, this makes sense. I realized that the root cause is already known and it's not possible to cause issue else where.

petern48 · 2025-12-13T20:31:45Z

datafusion/sqllogictest/test_files/aggregate.slt

+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING


I wasn't familiar with these before, but this was a great idea! It helped me find and understand a bug.

petern48 · 2025-12-13T20:35:35Z

datafusion/core/tests/dataframe/mod.rs

+    | -85         | -101     | 14              | -12           | -12    | 83  | -101 | 4  | -54  |
+    | -85         | -101     | 17              | -25           | -25    | 83  | -101 | 5  | -31  |


I found that this test was returning incorrect results due to the bug I explained in another comment, instead of raising an error. The results here were fixed by updating evaluate() to pass a &mut instead of consuming the state with std::mem::take().

Jefffrey

Looks good to me, just one minor question on some of the updated test results

Jefffrey · 2025-12-14T15:08:25Z

datafusion/core/tests/dataframe/mod.rs

+    | -85         | -48      | 6               | -35           | -36    | 83  | -85  | 2  | -43  |
+    | -85         | -5       | 4               | -37           | -40    | -5  | -85  | 1  | 83   |
+    | -85         | -54      | 15              | -17           | -18    | 83  | -101 | 4  | -38  |
+    | -85         | -56      | 2               | -70           | 57     | -56 | -85  | 1  | -25  |


I find this interesting, how we have -70 for the approx median but 57 for median 🤔

Great catch. I looked into it, and it seems like it's wrapping around due to integer overflow while taking the average of the middle two values (since the count is even).

low: [-85], high: -56, median: 57 datatype: Int8

-85 + -56 = -141 -> wraparound to 115
Then 115 / 2 -> 57.5 -> 57 (truncated due to integer type)

What's our desired behavior in this case? We could promote to a larger datatype to perform the calculation. Also is it intentional to return the value as a truncated integer instead of a float?

Regarding overflow, perhaps we should raise a separate issue to discuss/track this, as it does seem like incorrect behaviour.

We could do similar for the truncated integer behaviour; there was a recent issue asking about this for reference: #18867 (comment)

Filed: #19322

Jefffrey · 2025-12-14T15:18:24Z

datafusion/sqllogictest/test_files/aggregate.slt

+# median_non_sliding_window
+query ITRRRR
+SELECT
+    timestamp,
+    tags,
+    value,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
+    ) AS value_median_unbounded_preceding,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+    ) AS value_median_unbounded_both,


There is this quite informative comment which seems to explain why this is the case:

datafusion/datafusion/physical-expr/src/aggregate.rs

Lines 490 to 538 in befaf93

// Accumulators that have window frame startings different

// than `UNBOUNDED PRECEDING`, such as `1 PRECEDING`, need to

// implement retract_batch method in order to run correctly

// currently in DataFusion.

//

// If this `retract_batches` is not present, there is no way

// to calculate result correctly. For example, the query

//

// ```sql

// SELECT

// SUM(a) OVER(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS sum_a

// FROM

// t

// ```

//

// 1. First sum value will be the sum of rows between `[0, 1)`,

//

// 2. Second sum value will be the sum of rows between `[0, 2)`

//

// 3. Third sum value will be the sum of rows between `[1, 3)`, etc.

//

// Since the accumulator keeps the running sum:

//

// 1. First sum we add to the state sum value between `[0, 1)`

//

// 2. Second sum we add to the state sum value between `[1, 2)`

// (`[0, 1)` is already in the state sum, hence running sum will

// cover `[0, 2)` range)

//

// 3. Third sum we add to the state sum value between `[2, 3)`

// (`[0, 2)` is already in the state sum). Also we need to

// retract values between `[0, 1)` by this way we can obtain sum

// between [1, 3) which is indeed the appropriate range.

//

// When we use `UNBOUNDED PRECEDING` in the query starting

// index will always be 0 for the desired range, and hence the

// `retract_batch` method will not be called. In this case

// having retract_batch is not a requirement.

//

// This approach is a a bit different than window function

// approach. In window function (when they use a window frame)

// they get all the desired range during evaluation.

if !accumulator.supports_retract_batch() {

return not_impl_err!(

"Aggregate can not be used as a sliding accumulator because \

`retract_batch` is not implemented: {}",

self.name

);

}

Dandandan · 2025-12-14T20:02:26Z

datafusion/functions-aggregate/src/median.rs

+    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+        let values = values[0].as_primitive::<T>();
+        for v in values.iter().flatten() {
+            if let Some(idx) = self.all_values.iter().position(|x| *x == v) {


It seems this could be very slow?

Thanks, I improved it using a hashmap in 1b710cc

2010YOUY01 · 2025-12-15T04:04:34Z

datafusion/sqllogictest/test_files/aggregate.slt

+# median_non_sliding_window
+query ITRRRR
+SELECT
+    timestamp,
+    tags,
+    value,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
+    ) AS value_median_unbounded_preceding,
+    median(value) OVER (
+        PARTITION BY tags
+        ORDER BY timestamp
+        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+    ) AS value_median_unbounded_both,


I think we should file a ticket, the previous impl should be able to handle unbounded preceding as @Jefffrey explained, and the inconsistent results is likely to indicate a bug.

2010YOUY01 · 2025-12-15T04:05:39Z

datafusion/functions-aggregate/src/median.rs

    }
+
+    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+        let mut to_remove: HashMap<ScalarValue, usize> = HashMap::new();


This seems like a good optimization with minimal added complexity.

Implement 'retract_batch' for MedianAccumulator

57e768d

github-actions bot added sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Dec 11, 2025

petern48 commented Dec 11, 2025

View reviewed changes

datafusion/functions-aggregate/src/median.rs Outdated Show resolved Hide resolved

2010YOUY01 reviewed Dec 11, 2025

View reviewed changes

petern48 added 3 commits December 13, 2025 11:34

Pass &mut instead of cloning

81ced74

Add slt test with 'UNBOUNDED PRECEDING/FOLLOWING

5771f83

Fix expected results in mod.rs 'window_using_aggregates' and add test…

637c0b9

… for the bug in aggregate.slt

github-actions bot added the core Core DataFusion crate label Dec 13, 2025

petern48 added 2 commits December 13, 2025 12:15

cargo fmt

839e54e

Remove redundant test

768e95e

petern48 commented Dec 13, 2025

View reviewed changes

petern48 requested a review from 2010YOUY01 December 13, 2025 22:09

petern48 marked this pull request as ready for review December 13, 2025 22:09

petern48 changed the title ~~feat: Implement 'retract_batch' for MedianAccumulator~~ feat: Support sliding window queries for MedianAccumulator by implementing retract_batch Dec 13, 2025

Jefffrey approved these changes Dec 14, 2025

View reviewed changes

Dandandan reviewed Dec 14, 2025

View reviewed changes

Speed up retract_batch using a hash map

1b710cc

2010YOUY01 approved these changes Dec 15, 2025

View reviewed changes

petern48 mentioned this pull request Dec 15, 2025

bug: Median() encountered integer overflow #19322

Closed

2010YOUY01 added this pull request to the merge queue Dec 16, 2025

Merged via the queue into apache:main with commit 933657e Dec 16, 2025
28 checks passed

petern48 deleted the median_retract_batch branch December 16, 2025 15:25

petern48 mentioned this pull request Dec 29, 2025

bug: Median() truncates integers #19536

Open

This was referenced Jan 2, 2026

Accumulators which don't implement retract_batch can still exhibit buggy behaviour #19612

Closed

fix(accumulators): preserve state in evaluate() for window frame queries #19618

Merged

implement var distinct #19706

Merged

	// Accumulators that have window frame startings different
	// than `UNBOUNDED PRECEDING`, such as `1 PRECEDING`, need to
	// implement retract_batch method in order to run correctly
	// currently in DataFusion.
	//
	// If this `retract_batches` is not present, there is no way
	// to calculate result correctly. For example, the query
	//
	// ```sql
	// SELECT
	// SUM(a) OVER(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS sum_a
	// FROM
	// t
	// ```
	//
	// 1. First sum value will be the sum of rows between `[0, 1)`,
	//
	// 2. Second sum value will be the sum of rows between `[0, 2)`
	//
	// 3. Third sum value will be the sum of rows between `[1, 3)`, etc.
	//
	// Since the accumulator keeps the running sum:
	//
	// 1. First sum we add to the state sum value between `[0, 1)`
	//
	// 2. Second sum we add to the state sum value between `[1, 2)`
	// (`[0, 1)` is already in the state sum, hence running sum will
	// cover `[0, 2)` range)
	//
	// 3. Third sum we add to the state sum value between `[2, 3)`
	// (`[0, 2)` is already in the state sum). Also we need to
	// retract values between `[0, 1)` by this way we can obtain sum
	// between [1, 3) which is indeed the appropriate range.
	//
	// When we use `UNBOUNDED PRECEDING` in the query starting
	// index will always be 0 for the desired range, and hence the
	// `retract_batch` method will not be called. In this case
	// having retract_batch is not a requirement.
	//
	// This approach is a a bit different than window function
	// approach. In window function (when they use a window frame)
	// they get all the desired range during evaluation.
	if !accumulator.supports_retract_batch() {
	return not_impl_err!(
	"Aggregate can not be used as a sliding accumulator because \
	`retract_batch` is not implemented: {}",
	self.name
	);
	}

		\| -85 \| -101 \| 14 \| -12 \| -12 \| 83 \| -101 \| 4 \| -54 \|
		\| -85 \| -101 \| 17 \| -25 \| -25 \| 83 \| -101 \| 5 \| -31 \|

Conversation

petern48 commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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?

Uh oh!

Uh oh!

2010YOUY01 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

petern48 commented Dec 11, 2025 •

edited

Loading