kosiew commented on code in PR #26417:
URL: https://github.com/apache/datafusion/pull/26417#discussion_r4322089002
##########
datafusion/core/tests/physical_optimizer/window_topn.rs:
##########
@@ -458,6 +461,137 @@ fn two_window_exprs_keep_the_window_node() -> Result<()> {
Ok(())
}
+/// `ROW_NUMBER() AS rn, RANK() AS rk` over `PARTITION BY pk ORDER BY val`.
+/// Output: `[pk@0, val@1, rn@2, rk@3]`.
+fn build_rn_rk_window() -> Result<Arc<dyn ExecutionPlan>> {
+ let s = schema();
+ let input: Arc<dyn ExecutionPlan> =
Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+ let partition_by = vec![col("pk", &s)?];
+ let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()];
+ let window_expr = |udwf: Arc<datafusion_expr::WindowUDF>, name: &str| {
+ Ok::<_,
datafusion_common::DataFusionError>(Arc::new(StandardWindowExpr::new(
+ create_udwf_window_expr(&udwf, &[], &s, name.to_string(), false)?,
+ &partition_by,
+ &order_by,
+ Arc::new(WindowFrame::new_bounds(
+ WindowFrameUnits::Rows,
+ WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+ WindowFrameBound::CurrentRow,
+ )),
+ )))
+ };
+ Ok(Arc::new(BoundedWindowAggExec::try_new(
+ vec![
+ window_expr(row_number_udwf(), "rn")?,
+ window_expr(rank_udwf(), "rk")?,
+ ],
+ input,
+ InputOrderMode::Sorted,
+ true,
+ )?))
+}
+
+/// A `ProjectionExec` passing through `input`'s columns at `indices`.
+fn project(
+ input: Arc<dyn ExecutionPlan>,
+ indices: &[usize],
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = input.schema();
+ let exprs: Vec<(Arc<dyn datafusion_physical_expr::PhysicalExpr>, String)>
= indices
+ .iter()
+ .map(|&i| {
+ let name = schema.field(i).name();
+ (
+ Arc::new(Column::new(name, i))
+ as Arc<dyn datafusion_physical_expr::PhysicalExpr>,
+ name.to_string(),
+ )
+ })
+ .collect();
+ Ok(Arc::new(ProjectionExec::try_new(exprs, input)?))
+}
+
+/// `FilterExec(<column idx of input> <= 3)`, with the literal in the column's
type.
+fn filter_lteq_3(
+ input: Arc<dyn ExecutionPlan>,
+ idx: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = input.schema();
+ let field = schema.field(idx);
+ let predicate = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new(field.name(), idx)),
+ Operator::LtEq,
+ lit(ScalarValue::UInt64(Some(3)).cast_to(field.data_type())?),
+ ));
+ Ok(Arc::new(FilterExec::try_new(predicate, input)?))
+}
+
+/// The ranking function the rewrite chose, or `None` if it did not fire.
+fn rewritten_fn(plan: Arc<dyn ExecutionPlan>) -> Result<Option<WindowFnKind>> {
+ let optimized = optimize(plan)?;
+ Ok(find_partitioned_topk(&optimized).map(|p| {
+ p.downcast_ref::<PartitionedTopKExec>()
+ .expect("find_partitioned_topk returns a PartitionedTopKExec")
+ .fn_kind()
+ }))
+}
+
+/// Stacked projections are resolved top-down: each one maps an index into its
+/// output to one into its input. Bottom-up, `rk@0` would follow `pk@0` and
then
+/// `rk@2` to `rn` in the window's output.
+#[test]
+fn stacked_projections_resolve_top_down() -> Result<()> {
Review Comment:
Could we add a negative test using `project(build_rn_rk_window()?, &[2, 3,
0])` followed by `filter_lteq_3(..., 2)` and verify the plan stays unchanged?
This checks that a reordered input column is correctly mapped back to `pk@0`
instead of being mistaken for `rn`.
##########
datafusion/sqllogictest/test_files/window_topn.slt:
##########
@@ -1541,3 +1541,162 @@ set datafusion.execution.batch_size = 8192;
statement ok
set datafusion.optimizer.enable_window_topn = false;
+
+# The filter's column index is resolved through the projection between the
+# filter and the window. Here the outer query drops `v`, so the projection
+# shifts every later column down by one; with two ranking functions, using the
+# filter's index as if it were the window's picks the other function.
+statement ok
+CREATE TABLE window_topn_shift (pk INT, v INT, id INT) AS VALUES
+ (1, 0, 1), (1, 0, 2), (1, 1, 3), (2, 0, 4), (2, 0, 5), (2, 1, 6);
+
+statement ok
+SET datafusion.optimizer.enable_window_topn = true;
+
+statement ok
+SET datafusion.explain.physical_plan_only = true;
+
+# Filter on rank(): both tied rows per partition are kept.
+query IIII
+SELECT pk, id, rn, rk FROM (
Review Comment:
Since `ORDER BY v` doesn't define the order of tied rows, these tests
shouldn't assume which `id` gets each `row_number` or which peer survives `rn
<= 1`. Could we check the retained IDs and row numbers independently, and allow
either tied peer for `rn <= 1`? Please keep the ties in the window ordering so
the regression still catches the wrong-function bug.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
