Include data types in logical plans of inferred prepare statements (apache#16019)

* draft commit to rolledback changes on function naming and include prepare clause on the infer types tests

* include data types in plan when it is not included in the prepare statement

* fix: prepare statement error

* Update datafusion/sql/src/statement.rs

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

* remove infer types from prepare statement

the infer data type changes in statement will be introduced in a new PR

* fix to show correct output message

* include data types on logical plans of prepare statements without explicit type declaration

* fix using clippy sugestions

* explicitly get the data types using the placeholder id to avoid sorting

* Restore the original tests too

* update set data type routine to be more rust idiomatic

Co-authored-by: Tommy shu <qstommyshu@gmail.com>

* update set datatype routine

* fix formatting in sql_integration

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
Co-authored-by: Tommy shu <qstommyshu@gmail.com>

Author: 3 people

datafusion/sql/src/statement.rs

@@ -696,7 +696,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 statement,
             } => {
                 // Convert parser data types to DataFusion data types
-                let data_types: Vec<DataType> = data_types
+                let mut data_types: Vec<DataType> = data_types
                     .into_iter()
                     .map(|t| self.convert_data_type(&t))
                     .collect::<Result<_>>()?;
@@ -710,6 +710,19 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     *statement,
                     &mut planner_context,
                 )?;
+
+                if data_types.is_empty() {
+                    let map_types = plan.get_parameter_types()?;
+                    let param_types: Vec<_> = (1..=map_types.len())
+                        .filter_map(|i| {
+                            let key = format!("${i}");
+                            map_types.get(&key).and_then(|opt| opt.clone())
+                        })
+                        .collect();
+                    data_types.extend(param_types.iter().cloned());
+                    planner_context.with_prepare_param_data_types(param_types);
+                }
+
                 Ok(LogicalPlan::Statement(PlanStatement::Prepare(Prepare {
                     name: ident_to_string(&name),
                     data_types,

datafusion/sql/tests/sql_integration.rs

@@ -4650,7 +4650,7 @@ fn test_prepare_statement_infer_types_from_join() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [Int32] 
       Projection: person.id, orders.order_id
         Inner Join:  Filter: person.id = orders.customer_id AND person.age = $1
           TableScan: person
@@ -4661,6 +4661,20 @@ fn test_prepare_statement_infer_types_from_join() {
     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);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::Int32(Some(10))];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+
+    assert_snapshot!(
+        plan_with_params,
+        @r"
+    Projection: person.id, orders.order_id
+      Inner Join:  Filter: person.id = orders.customer_id AND person.age = Int32(10)
+        TableScan: person
+        TableScan: orders
+    "
+    );
 }
 
 #[test]
@@ -4701,7 +4715,7 @@ fn test_prepare_statement_infer_types_from_predicate() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [Int32] 
       Projection: person.id, person.age
         Filter: person.age = $1
           TableScan: person
@@ -4711,6 +4725,19 @@ fn test_prepare_statement_infer_types_from_predicate() {
     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);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::Int32(Some(10))];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+
+    assert_snapshot!(
+        plan_with_params,
+        @r"
+    Projection: person.id, person.age
+      Filter: person.age = Int32(10)
+        TableScan: person
+    "
+    );
 }
 
 #[test]
@@ -4756,7 +4783,7 @@ fn test_prepare_statement_infer_types_from_between_predicate() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [Int32, Int32] 
       Projection: person.id, person.age
         Filter: person.age BETWEEN $1 AND $2
           TableScan: person
@@ -4769,6 +4796,19 @@ fn test_prepare_statement_infer_types_from_between_predicate() {
         ("$2".to_string(), Some(DataType::Int32)),
     ]);
     assert_eq!(actual_types, expected_types);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(30))];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+
+    assert_snapshot!(
+        plan_with_params,
+        @r"
+    Projection: person.id, person.age
+      Filter: person.age BETWEEN Int32(10) AND Int32(30)
+        TableScan: person
+    "
+    );
 }
 
 #[test]
@@ -4821,7 +4861,7 @@ fn test_prepare_statement_infer_types_subquery() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [UInt32] 
       Projection: person.id, person.age
         Filter: person.age = (<subquery>)
           Subquery:
@@ -4836,6 +4876,24 @@ fn test_prepare_statement_infer_types_subquery() {
     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);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::UInt32(Some(10))];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+
+    assert_snapshot!(
+        plan_with_params,
+        @r"
+    Projection: person.id, person.age
+      Filter: person.age = (<subquery>)
+        Subquery:
+          Projection: max(person.age)
+            Aggregate: groupBy=[[]], aggr=[[max(person.age)]]
+              Filter: person.id = UInt32(10)
+                TableScan: person
+        TableScan: person
+        "
+    );
 }
 
 #[test]
@@ -4883,7 +4941,7 @@ fn test_prepare_statement_update_infer() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [Int32, UInt32] 
       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
@@ -4897,6 +4955,20 @@ fn test_prepare_statement_update_infer() {
         ("$2".to_string(), Some(DataType::UInt32)),
     ]);
     assert_eq!(actual_types, expected_types);
+
+    // replace params with values
+    let param_values = vec![ScalarValue::Int32(Some(42)), ScalarValue::UInt32(Some(1))];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+
+    assert_snapshot!(
+        plan_with_params,
+        @r"
+    Dml: op=[Update] table=[person]
+      Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, Int32(42) AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
+        Filter: person.id = UInt32(1)
+          TableScan: person
+        "
+    );
 }
 
 #[test]
@@ -4944,7 +5016,7 @@ fn test_prepare_statement_insert_infer() {
     assert_snapshot!(
         plan,
         @r#"
-    Prepare: "my_plan" [] 
+    Prepare: "my_plan" [UInt32, Utf8, Utf8] 
       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)
@@ -4958,6 +5030,22 @@ fn test_prepare_statement_insert_infer() {
         ("$3".to_string(), Some(DataType::Utf8)),
     ]);
     assert_eq!(actual_types, expected_types);
+
+    // replace params with values
+    let param_values = vec![
+        ScalarValue::UInt32(Some(1)),
+        ScalarValue::from("Alan"),
+        ScalarValue::from("Turing"),
+    ];
+    let plan_with_params = plan.with_param_values(param_values).unwrap();
+    assert_snapshot!(
+        plan_with_params,
+        @r#"
+    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: (UInt32(1) AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
+    "#
+    );
 }
 
 #[test]

datafusion/sqllogictest/test_files/prepare.slt

@@ -92,7 +92,7 @@ DEALLOCATE my_plan
 statement ok
 PREPARE my_plan AS SELECT * FROM person WHERE id < $1;
 
-statement error No value found for placeholder with id \$1
+statement error Prepared statement 'my_plan' expects 1 parameters, but 0 provided
 EXECUTE my_plan
 
 statement ok