fix: planning of prepare statement with limit clause (#13088)

* fix: planning of prepare statement with limit clause

* Improve test

Author: jonahgao

datafusion/sql/src/query.rs

@@ -53,7 +53,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 // so we need to process `SELECT` and `ORDER BY` together.
                 let oby_exprs = to_order_by_exprs(query.order_by)?;
                 let plan = self.select_to_plan(*select, oby_exprs, planner_context)?;
-                let plan = self.limit(plan, query.offset, query.limit)?;
+                let plan =
+                    self.limit(plan, query.offset, query.limit, planner_context)?;
                 // Process the `SELECT INTO` after `LIMIT`.
                 self.select_into(plan, select_into)
             }
@@ -68,7 +69,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                     None,
                 )?;
                 let plan = self.order_by(plan, order_by_rex)?;
-                self.limit(plan, query.offset, query.limit)
+                self.limit(plan, query.offset, query.limit, planner_context)
             }
         }
     }
@@ -79,6 +80,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         input: LogicalPlan,
         skip: Option<SQLOffset>,
         fetch: Option<SQLExpr>,
+        planner_context: &mut PlannerContext,
     ) -> Result<LogicalPlan> {
         if skip.is_none() && fetch.is_none() {
             return Ok(input);
@@ -88,10 +90,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         let empty_schema = DFSchema::empty();
 
         let skip = skip
-            .map(|o| self.sql_to_expr(o.value, &empty_schema, &mut PlannerContext::new()))
+            .map(|o| self.sql_to_expr(o.value, &empty_schema, planner_context))
             .transpose()?;
         let fetch = fetch
-            .map(|e| self.sql_to_expr(e, &empty_schema, &mut PlannerContext::new()))
+            .map(|e| self.sql_to_expr(e, &empty_schema, planner_context))
             .transpose()?;
         LogicalPlanBuilder::from(input)
             .limit_by_expr(skip, fetch)?

datafusion/sql/tests/sql_integration.rs

@@ -4209,6 +4209,29 @@ fn test_prepare_statement_to_plan_having() {
     prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan);
 }
 
+#[test]
+fn test_prepare_statement_to_plan_limit() {
+    let sql = "PREPARE my_plan(BIGINT, BIGINT) AS
+        SELECT id FROM person \
+        OFFSET $1 LIMIT $2";
+
+    let expected_plan = "Prepare: \"my_plan\" [Int64, Int64] \
+        \n  Limit: skip=$1, fetch=$2\
+        \n    Projection: person.id\
+        \n      TableScan: person";
+
+    let expected_dt = "[Int64, Int64]";
+
+    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::Int64(Some(10)), ScalarValue::Int64(Some(200))];
+    let expected_plan = "Limit: skip=10, fetch=200\
+        \n  Projection: person.id\
+        \n    TableScan: person";
+    prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan);
+}
+
 #[test]
 fn test_prepare_statement_to_plan_value_list() {
     let sql = "PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);";