SubhamSinghal commented on code in PR #26417:
URL: https://github.com/apache/datafusion/pull/26417#discussion_r4369918098
##########
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:
addressed in 0ca88916b76136fc7d1cb542996397368c2fb29e
##########
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:
Addressed in 0ca88916b76136fc7d1cb542996397368c2fb29e
--
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]
