Skip to content

Commit 79ecf94

Browse files
mustafasrepoozankabak
mustafasrepo
and
ozankabak
authored
Linear search support for Window Group queries (#5286)
* add naive linear search * Add last range to decrease search size * minor changes * add low, high arguments * Go back to old API, improve comments, refactors * Linear Groups implementation * Resolve linter errors * remove old unit tests * simplifications * Add unit tests * Remove sort options from GROUPS calculations, various code simplifications and comment clarifications * New TODOs to fix * Address reviews * Fix error * Prehandle range current row and unbounded following case * Fix error * Move a check from execution to planning, reduce code duplication * Incorporate review suggestion (with cargo fmt fix) --------- Co-authored-by: Mehmet Ozan Kabak <[email protected]>
1 parent ded897e commit 79ecf94

9 files changed

Lines changed: 420 additions & 763 deletions

File tree

‎datafusion/core/src/physical_plan/planner.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1594,7 +1594,7 @@ pub fn create_window_expr_with_name(
15941594
})
15951595
.collect::<Result<Vec<_>>>()?;
15961596
if !is_window_valid(window_frame) {
1597-
return Err(DataFusionError::Execution(format!(
1597+
return Err(DataFusionError::Plan(format!(
15981598
"Invalid window frame: start bound ({}) cannot be larger than end bound ({})",
15991599
window_frame.start_bound, window_frame.end_bound
16001600
)));

‎datafusion/core/tests/sql/window.rs‎

Lines changed: 54 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -982,19 +982,20 @@ async fn window_frame_groups_multiple_order_columns() -> Result<()> {
982982
async fn window_frame_groups_without_order_by() -> Result<()> {
983983
let ctx = SessionContext::new();
984984
register_aggregate_csv(&ctx).await?;
985-
// execute the query
986-
let df = ctx
985+
// Try executing an erroneous query (the ORDER BY clause is missing in the
986+
// window frame):
987+
let err = ctx
987988
.sql(
988989
"SELECT
989990
SUM(c4) OVER(PARTITION BY c2 GROUPS BETWEEN 1 PRECEDING AND 1 FOLLOWING)
990991
FROM aggregate_test_100
991992
ORDER BY c9;",
992993
)
993-
.await?;
994-
let err = df.collect().await.unwrap_err();
994+
.await
995+
.unwrap_err();
995996
assert_contains!(
996997
err.to_string(),
997-
"Execution error: GROUPS mode requires an ORDER BY clause".to_owned()
998+
"Error during planning: GROUPS mode requires an ORDER BY clause".to_owned()
998999
);
9991000
Ok(())
10001001
}
@@ -1034,7 +1035,7 @@ async fn window_frame_creation() -> Result<()> {
10341035
let results = df.collect().await;
10351036
assert_eq!(
10361037
results.err().unwrap().to_string(),
1037-
"Execution error: Invalid window frame: start bound (1 PRECEDING) cannot be larger than end bound (2 PRECEDING)"
1038+
"Error during planning: Invalid window frame: start bound (1 PRECEDING) cannot be larger than end bound (2 PRECEDING)"
10381039
);
10391040

10401041
let df = ctx
@@ -1047,7 +1048,20 @@ async fn window_frame_creation() -> Result<()> {
10471048
let results = df.collect().await;
10481049
assert_eq!(
10491050
results.err().unwrap().to_string(),
1050-
"Execution error: Invalid window frame: start bound (2 FOLLOWING) cannot be larger than end bound (1 FOLLOWING)"
1051+
"Error during planning: Invalid window frame: start bound (2 FOLLOWING) cannot be larger than end bound (1 FOLLOWING)"
1052+
);
1053+
1054+
let err = ctx
1055+
.sql(
1056+
"SELECT
1057+
COUNT(c1) OVER(GROUPS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)
1058+
FROM aggregate_test_100;",
1059+
)
1060+
.await
1061+
.unwrap_err();
1062+
assert_contains!(
1063+
err.to_string(),
1064+
"Error during planning: GROUPS mode requires an ORDER BY clause"
10511065
);
10521066

10531067
Ok(())
@@ -1123,6 +1137,39 @@ async fn test_window_row_number_aggregate() -> Result<()> {
11231137
Ok(())
11241138
}
11251139

1140+
#[tokio::test]
1141+
async fn test_window_range_equivalent_frames() -> Result<()> {
1142+
let config = SessionConfig::new();
1143+
let ctx = SessionContext::with_config(config);
1144+
register_aggregate_csv(&ctx).await?;
1145+
let sql = "SELECT
1146+
c9,
1147+
COUNT(*) OVER(ORDER BY c9, c1 RANGE BETWEEN CURRENT ROW AND CURRENT ROW) AS cnt1,
1148+
COUNT(*) OVER(ORDER BY c9, c1 RANGE UNBOUNDED PRECEDING) AS cnt2,
1149+
COUNT(*) OVER(ORDER BY c9, c1 RANGE CURRENT ROW) AS cnt3,
1150+
COUNT(*) OVER(RANGE BETWEEN CURRENT ROW AND CURRENT ROW) AS cnt4,
1151+
COUNT(*) OVER(RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS cnt5,
1152+
COUNT(*) OVER(RANGE BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING) AS cnt6
1153+
FROM aggregate_test_100
1154+
ORDER BY c9
1155+
LIMIT 5";
1156+
1157+
let actual = execute_to_batches(&ctx, sql).await;
1158+
let expected = vec![
1159+
"+-----------+------+------+------+------+------+------+",
1160+
"| c9 | cnt1 | cnt2 | cnt3 | cnt4 | cnt5 | cnt6 |",
1161+
"+-----------+------+------+------+------+------+------+",
1162+
"| 28774375 | 1 | 1 | 1 | 100 | 100 | 100 |",
1163+
"| 63044568 | 1 | 2 | 1 | 100 | 100 | 100 |",
1164+
"| 141047417 | 1 | 3 | 1 | 100 | 100 | 100 |",
1165+
"| 141680161 | 1 | 4 | 1 | 100 | 100 | 100 |",
1166+
"| 145294611 | 1 | 5 | 1 | 100 | 100 | 100 |",
1167+
"+-----------+------+------+------+------+------+------+",
1168+
];
1169+
assert_batches_eq!(expected, &actual);
1170+
Ok(())
1171+
}
1172+
11261173
#[tokio::test]
11271174
async fn test_window_cume_dist() -> Result<()> {
11281175
let config = SessionConfig::new();

‎datafusion/expr/src/window_frame.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,36 @@ impl WindowFrame {
144144
}
145145
}
146146

147+
/// Construct equivalent explicit window frames for implicit corner cases.
148+
/// With this processing, we may assume in downstream code that RANGE/GROUPS
149+
/// frames contain an appropriate ORDER BY clause.
150+
pub fn regularize(mut frame: WindowFrame, order_bys: usize) -> Result<WindowFrame> {
151+
if frame.units == WindowFrameUnits::Range && order_bys != 1 {
152+
// Normally, RANGE frames require an ORDER BY clause with exactly one
153+
// column. However, an ORDER BY clause may be absent in two edge cases.
154+
if (frame.start_bound.is_unbounded()
155+
|| frame.start_bound == WindowFrameBound::CurrentRow)
156+
&& (frame.end_bound == WindowFrameBound::CurrentRow
157+
|| frame.end_bound.is_unbounded())
158+
{
159+
if order_bys == 0 {
160+
frame.units = WindowFrameUnits::Rows;
161+
frame.start_bound =
162+
WindowFrameBound::Preceding(ScalarValue::UInt64(None));
163+
frame.end_bound = WindowFrameBound::Following(ScalarValue::UInt64(None));
164+
}
165+
} else {
166+
return Err(DataFusionError::Plan(format!(
167+
"With window frame of type RANGE, the ORDER BY expression must be of length 1, got {}", order_bys)));
168+
}
169+
} else if frame.units == WindowFrameUnits::Groups && order_bys == 0 {
170+
return Err(DataFusionError::Plan(
171+
"GROUPS mode requires an ORDER BY clause".to_string(),
172+
));
173+
};
174+
Ok(frame)
175+
}
176+
147177
/// There are five ways to describe starting and ending frame boundaries:
148178
///
149179
/// 1. UNBOUNDED PRECEDING

‎datafusion/physical-expr/src/window/built_in.rs‎

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -104,17 +104,15 @@ impl WindowExpr for BuiltInWindowExpr {
104104
let mut row_wise_results = vec![];
105105

106106
let (values, order_bys) = self.get_values_orderbys(batch)?;
107-
let mut window_frame_ctx = WindowFrameContext::new(&self.window_frame);
108-
let range = Range { start: 0, end: 0 };
107+
let mut window_frame_ctx = WindowFrameContext::new(
108+
&self.window_frame,
109+
sort_options,
110+
Range { start: 0, end: 0 },
111+
);
109112
// We iterate on each row to calculate window frame range and and window function result
110113
for idx in 0..num_rows {
111-
let range = window_frame_ctx.calculate_range(
112-
&order_bys,
113-
&sort_options,
114-
num_rows,
115-
idx,
116-
&range,
117-
)?;
114+
let range =
115+
window_frame_ctx.calculate_range(&order_bys, num_rows, idx)?;
118116
let value = evaluator.evaluate_inside_range(&values, &range)?;
119117
row_wise_results.push(value);
120118
}
@@ -168,7 +166,13 @@ impl WindowExpr for BuiltInWindowExpr {
168166
// We iterate on each row to perform a running calculation.
169167
let record_batch = &partition_batch_state.record_batch;
170168
let num_rows = record_batch.num_rows();
171-
let mut window_frame_ctx = WindowFrameContext::new(&self.window_frame);
169+
let last_range = state.window_frame_range.clone();
170+
let mut window_frame_ctx = WindowFrameContext::new(
171+
&self.window_frame,
172+
sort_options.clone(),
173+
// Start search from the last range
174+
last_range,
175+
);
172176
let sort_partition_points = if evaluator.include_rank() {
173177
let columns = self.sort_columns(record_batch)?;
174178
self.evaluate_partition_points(num_rows, &columns)?
@@ -179,13 +183,7 @@ impl WindowExpr for BuiltInWindowExpr {
179183
let mut last_range = state.window_frame_range.clone();
180184
for idx in state.last_calculated_index..num_rows {
181185
state.window_frame_range = if self.expr.uses_window_frame() {
182-
window_frame_ctx.calculate_range(
183-
&order_bys,
184-
&sort_options,
185-
num_rows,
186-
idx,
187-
&state.window_frame_range,
188-
)
186+
window_frame_ctx.calculate_range(&order_bys, num_rows, idx)
189187
} else {
190188
evaluator.get_range(state, num_rows)
191189
}?;

‎datafusion/physical-expr/src/window/window_expr.rs‎

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -162,18 +162,10 @@ pub trait AggregateWindowExpr: WindowExpr {
162162

163163
/// Evaluates the window function against the batch.
164164
fn aggregate_evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
165-
let mut window_frame_ctx = WindowFrameContext::new(self.get_window_frame());
166165
let mut accumulator = self.get_accumulator()?;
167166
let mut last_range = Range { start: 0, end: 0 };
168167
let mut idx = 0;
169-
self.get_result_column(
170-
&mut accumulator,
171-
batch,
172-
&mut window_frame_ctx,
173-
&mut last_range,
174-
&mut idx,
175-
false,
176-
)
168+
self.get_result_column(&mut accumulator, batch, &mut last_range, &mut idx, false)
177169
}
178170

179171
/// Statefully evaluates the window function against the batch. Maintains
@@ -207,11 +199,9 @@ pub trait AggregateWindowExpr: WindowExpr {
207199
let mut state = &mut window_state.state;
208200

209201
let record_batch = &partition_batch_state.record_batch;
210-
let mut window_frame_ctx = WindowFrameContext::new(self.get_window_frame());
211202
let out_col = self.get_result_column(
212203
accumulator,
213204
record_batch,
214-
&mut window_frame_ctx,
215205
&mut state.window_frame_range,
216206
&mut state.last_calculated_index,
217207
!partition_batch_state.is_end,
@@ -230,7 +220,6 @@ pub trait AggregateWindowExpr: WindowExpr {
230220
&self,
231221
accumulator: &mut Box<dyn Accumulator>,
232222
record_batch: &RecordBatch,
233-
window_frame_ctx: &mut WindowFrameContext,
234223
last_range: &mut Range<usize>,
235224
idx: &mut usize,
236225
not_end: bool,
@@ -240,15 +229,15 @@ pub trait AggregateWindowExpr: WindowExpr {
240229
let length = values[0].len();
241230
let sort_options: Vec<SortOptions> =
242231
self.order_by().iter().map(|o| o.options).collect();
232+
let mut window_frame_ctx = WindowFrameContext::new(
233+
self.get_window_frame(),
234+
sort_options,
235+
// Start search from the last range
236+
last_range.clone(),
237+
);
243238
let mut row_wise_results: Vec<ScalarValue> = vec![];
244239
while *idx < length {
245-
let cur_range = window_frame_ctx.calculate_range(
246-
&order_bys,
247-
&sort_options,
248-
length,
249-
*idx,
250-
last_range,
251-
)?;
240+
let cur_range = window_frame_ctx.calculate_range(&order_bys, length, *idx)?;
252241
// Exit if the range extends all the way:
253242
if cur_range.end == length && not_end {
254243
break;

0 commit comments

Comments
 (0)