fix: unable to get field on rank limit when rule_eager_aggregation applied (#19007)

Author: KKould

src/query/service/src/physical_plans/physical_aggregate_partial.rs

@@ -211,24 +211,12 @@ impl IPhysicalPlan for AggregatePartial {
                 .cluster_with_partial(true, builder.ctx.get_cluster().nodes.len())
         };
 
-        // For rank limit, we can filter data using sort with rank before partial
-        if let Some(rank_limit) = &self.rank_limit {
-            let sort_desc = rank_limit
-                .0
-                .iter()
-                .map(|desc| {
-                    let offset = schema_before_group_by.index_of(&desc.order_by.to_string())?;
-                    Ok(SortColumnDescription {
-                        offset,
-                        asc: desc.asc,
-                        nulls_first: desc.nulls_first,
-                    })
-                })
-                .collect::<Result<Vec<_>>>()?;
-            let sort_desc: Arc<[_]> = sort_desc.into();
-
+        // For rank limit, we can filter data using sort with rank before partial.
+        if let Some((sort_desc, limit)) =
+            self.resolve_rank_limit_descriptions(&schema_before_group_by)
+        {
             builder.main_pipeline.add_transformer(|| {
-                TransformSortPartial::new(LimitType::LimitRank(rank_limit.1), sort_desc.clone())
+                TransformSortPartial::new(LimitType::LimitRank(limit), sort_desc.clone())
             });
         }
 
@@ -283,3 +271,28 @@ impl IPhysicalPlan for AggregatePartial {
         Ok(())
     }
 }
+
+impl AggregatePartial {
+    fn resolve_rank_limit_descriptions(
+        &self,
+        schema_before_group_by: &DataSchemaRef,
+    ) -> Option<(Arc<[SortColumnDescription]>, usize)> {
+        let (sort_descs, limit) = self.rank_limit.as_ref()?;
+        let mut resolved = Vec::with_capacity(sort_descs.len());
+        for desc in sort_descs {
+            let field_name = desc.order_by.to_string();
+            let offset = match schema_before_group_by.index_of(&field_name) {
+                Ok(offset) => offset,
+                Err(_) => {
+                    return None;
+                }
+            };
+            resolved.push(SortColumnDescription {
+                offset,
+                asc: desc.asc,
+                nulls_first: desc.nulls_first,
+            });
+        }
+        Some((resolved.into(), *limit))
+    }
+}

tests/suites/0_stateless/20+_others/20_0027_rank_limit_eager.result

@@ -0,0 +1,5 @@
+0	0
+1	1
+2	2
+3	3
+4	4

tests/suites/0_stateless/20+_others/20_0027_rank_limit_eager.sql

@@ -0,0 +1,12 @@
+SELECT
+    b.number AS bn,
+    SUM(a.number) AS s
+FROM
+    numbers(100) AS a
+    JOIN numbers(100) AS b ON a.number = b.number
+GROUP BY
+    b.number
+ORDER BY
+    bn
+LIMIT
+    5;