apache · alamb · May 15, 2025 · Apr 17, 2025 · May 1, 2025 · May 1, 2025
diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs
@@ -4641,6 +4641,28 @@ fn test_infer_types_from_join() {
     );
     /// Verify parameter list length and type 
     pub fn verify(&self, expect: &[DataType]) -> Result<()> { 
         match self { 
             ParamValues::List(list) => { 
                 // Verify if the number of params matches the number of values 
                 if expect.len() != list.len() { 
                     return _plan_err!( 
                         "Expected {} parameters, got {}", 
                         expect.len(), 
                         list.len() 
                     ); 
                 } 
     /// Verify parameter list length and type 
     pub fn verify(&self, expect: &[DataType]) -> Result<()> { 
         match self { 
             ParamValues::List(list) => { 
                 // Verify if the number of params matches the number of values 
                 if expect.len() != list.len() { 
                     return _plan_err!( 
                         "Expected {} parameters, got {}", 
                         expect.len(), 
                         list.len() 
                     ); 
                 } 
 }
 
+#[test]
+fn test_prepare_statement_infer_types_from_join() {
+    let sql =
+        "PREPARE my_plan AS SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1";
+
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Projection: person.id, orders.order_id
+        Inner Join:  Filter: person.id = orders.customer_id AND person.age = $1
+          TableScan: person
+          TableScan: orders
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_infer_types_from_predicate() {
     let sql = "SELECT id, age FROM person WHERE age = $1";
@@ -4672,6 +4694,25 @@ fn test_infer_types_from_predicate() {
     );
 }
 
+#[test]
+fn test_prepare_statement_infer_types_from_predicate() {
+    let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age = $1";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Projection: person.id, person.age
+        Filter: person.age = $1
+          TableScan: person
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_infer_types_from_between_predicate() {
     let sql = "SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
@@ -4707,6 +4748,29 @@ fn test_infer_types_from_between_predicate() {
     );
 }
 
+#[test]
+fn test_prepare_statement_infer_types_from_between_predicate() {
+    let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
+
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Projection: person.id, person.age
+        Filter: person.age BETWEEN $1 AND $2
+          TableScan: person
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([
+        ("$1".to_string(), Some(DataType::Int32)),
+        ("$2".to_string(), Some(DataType::Int32)),
+    ]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_infer_types_subquery() {
     let sql = "SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)";
@@ -4749,6 +4813,31 @@ fn test_infer_types_subquery() {
     );
 }
 
+#[test]
+fn test_prepare_statement_infer_types_subquery() {
+    let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)";
+
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Projection: person.id, person.age
+        Filter: person.age = (<subquery>)
+          Subquery:
+            Projection: max(person.age)
+              Aggregate: groupBy=[[]], aggr=[[max(person.age)]]
+                Filter: person.id = $1
+                  TableScan: person
+          TableScan: person
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([("$1".to_string(), Some(DataType::UInt32))]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_update_infer() {
     let sql = "update person set age=$1 where id=$2";
@@ -4786,6 +4875,30 @@ fn test_update_infer() {
     );
 }
 
+#[test]
+fn test_prepare_statement_update_infer() {
+    let sql = "PREPARE my_plan AS update person set age=$1 where id=$2";
+
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Dml: op=[Update] table=[person]
+        Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, $1 AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
+          Filter: person.id = $2
+            TableScan: person
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([
+        ("$1".to_string(), Some(DataType::Int32)),
+        ("$2".to_string(), Some(DataType::UInt32)),
+    ]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_insert_infer() {
     let sql = "insert into person (id, first_name, last_name) values ($1, $2, $3)";
@@ -4824,6 +4937,29 @@ fn test_insert_infer() {
     );
 }
 
+#[test]
+fn test_prepare_statement_insert_infer() {
+    let sql = "PREPARE my_plan AS insert into person (id, first_name, last_name) values ($1, $2, $3)";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Dml: op=[Insert Into] table=[person]
+        Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
+          Values: ($1, $2, $3)
+    "#
+    );
+
+    let actual_types = plan.get_parameter_types().unwrap();
+    let expected_types = HashMap::from([
+        ("$1".to_string(), Some(DataType::UInt32)),
+        ("$2".to_string(), Some(DataType::Utf8)),
+        ("$3".to_string(), Some(DataType::Utf8)),
+    ]);
+    assert_eq!(actual_types, expected_types);
+}
+
 #[test]
 fn test_prepare_statement_to_plan_one_param() {
     let sql = "PREPARE my_plan(INT) AS SELECT id, age  FROM person WHERE age = $1";