Resolve comment

Author: Subham Singhal

datafusion/common/src/config.rs

@@ -1091,7 +1091,7 @@ config_namespace! {
         /// Filter(rn<=K) → Window(ROW_NUMBER) → Sort patterns with a
         /// PartitionedTopKExec that maintains per-partition heaps, avoiding
         /// a full sort of the input.
-        pub enable_window_topn: bool, default = true
+        pub enable_window_topn: bool, default = false
 
         /// When set to true, the optimizer will push TopK (Sort with fetch)
         /// below hash repartition when the partition key is a prefix of the

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

@@ -232,6 +232,7 @@ impl PartitionedTopKExec {
     /// The output is sorted by `sort_exprs` (partition keys then order keys),
     /// uses the same partitioning as the input, emits all output at once
     /// (`EmissionType::Final`), and is bounded.
+
     fn compute_properties(
         input: &Arc<dyn ExecutionPlan>,
         sort_exprs: LexOrdering,
@@ -269,8 +270,17 @@ impl DisplayAs for PartitionedTopKExec {
                 )
             }
             DisplayFormatType::TreeRender => {
+                let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len]
+                    .iter()
+                    .map(|e| format!("{}", e.expr))
+                    .collect();
+                let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..]
+                    .iter()
+                    .map(|e| format!("{e}"))
+                    .collect();
                 writeln!(f, "fetch={}", self.fetch)?;
-                writeln!(f, "{}", self.expr)
+                writeln!(f, "partition=[{}]", partition_exprs.join(", "))?;
+                writeln!(f, "order=[{}]", order_exprs.join(", "))
             }
         }
     }

datafusion/sqllogictest/test_files/information_schema.slt

@@ -309,6 +309,7 @@ datafusion.optimizer.enable_topk_aggregation true
 datafusion.optimizer.enable_topk_dynamic_filter_pushdown true
 datafusion.optimizer.enable_topk_repartition true
 datafusion.optimizer.enable_window_limits true
+datafusion.optimizer.enable_window_topn false
 datafusion.optimizer.expand_views_at_output false
 datafusion.optimizer.filter_null_join_keys false
 datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values 150
@@ -453,6 +454,7 @@ datafusion.optimizer.enable_topk_aggregation true When set to true, the optimize
 datafusion.optimizer.enable_topk_dynamic_filter_pushdown true When set to true, the optimizer will attempt to push down TopK dynamic filters into the file scan phase.
 datafusion.optimizer.enable_topk_repartition true When set to true, the optimizer will push TopK (Sort with fetch) below hash repartition when the partition key is a prefix of the sort key, reducing data volume before the shuffle.
 datafusion.optimizer.enable_window_limits true When set to true, the optimizer will attempt to push limit operations past window functions, if possible
+datafusion.optimizer.enable_window_topn false When set to true, the optimizer will replace Filter(rn<=K) → Window(ROW_NUMBER) → Sort patterns with a PartitionedTopKExec that maintains per-partition heaps, avoiding a full sort of the input.
 datafusion.optimizer.expand_views_at_output false When set to true, if the returned type is a view type then the output will be coerced to a non-view. Coerces `Utf8View` to `LargeUtf8`, and `BinaryView` to `LargeBinary`.
 datafusion.optimizer.filter_null_join_keys false When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down.
 datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values 150 Maximum number of distinct values (rows) in the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides with more rows than this will use hash table lookups instead. Set to 0 to always use hash table lookups. This provides an additional limit beyond `hash_join_inlist_pushdown_max_size` to prevent very large IN lists that might not provide much benefit over hash table lookups. This uses the deduplicated row count once the build side has been evaluated. The default is 150 values per partition. This is inspired by Trino's `max-filter-keys-per-column` setting. See: <https://trino.io/docs/current/admin/dynamic-filtering.html#dynamic-filter-collection-thresholds>

datafusion/sqllogictest/test_files/window_topn.slt

@@ -30,6 +30,10 @@ CREATE TABLE window_topn_t (id INT, pk INT, val INT) AS VALUES
   (9, 3, 50),
   (10, 3, 75);
 
+# Enable the optimization for all tests
+statement ok
+SET datafusion.optimizer.enable_window_topn = true;
+
 # Test 1: Correct results for top-2 per partition
 query III rowsort
 SELECT id, pk, val FROM (
@@ -109,4 +113,4 @@ statement ok
 SET datafusion.optimizer.enable_window_topn = true;
 
 statement ok
-DROP TABLE window_topn_t;
+DROP TABLE window_topn_t;

docs/source/user-guide/configs.md

@@ -140,6 +140,7 @@ The following configuration settings are available:
 | datafusion.optimizer.enable_round_robin_repartition                     | true                      | When set to true, the physical plan optimizer will try to add round robin repartitioning to increase parallelism to leverage more CPU cores                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  |
 | datafusion.optimizer.enable_topk_aggregation                            | true                      | When set to true, the optimizer will attempt to perform limit operations during aggregations, if possible                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
 | datafusion.optimizer.enable_window_limits                               | true                      | When set to true, the optimizer will attempt to push limit operations past window functions, if possible                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
+| datafusion.optimizer.enable_window_topn                                 | true                      | When set to true, the optimizer will replace Filter(rn<=K) → Window(ROW_NUMBER) → Sort patterns with a PartitionedTopKExec that maintains per-partition heaps, avoiding a full sort of the input.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            |
 | datafusion.optimizer.enable_topk_repartition                            | true                      | When set to true, the optimizer will push TopK (Sort with fetch) below hash repartition when the partition key is a prefix of the sort key, reducing data volume before the shuffle.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                         |
 | datafusion.optimizer.enable_topk_dynamic_filter_pushdown                | true                      | When set to true, the optimizer will attempt to push down TopK dynamic filters into the file scan phase.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
 | datafusion.optimizer.enable_join_dynamic_filter_pushdown                | true                      | When set to true, the optimizer will attempt to push down Join dynamic filters into the file scan phase.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |