Support named query parameters

datafusion/common/src/param_value.rs

@@ -126,8 +126,24 @@ impl From<Vec<ScalarValue>> for ParamValues {
     }
 }
 
-impl From<HashMap<String, ScalarValue>> for ParamValues {
-    fn from(value: HashMap<String, ScalarValue>) -> Self {
+impl<K> From<Vec<(K, ScalarValue)>> for ParamValues
+where
+    K: Into<String>,
+{
+    fn from(value: Vec<(K, ScalarValue)>) -> Self {
+        let value: HashMap<String, ScalarValue> =
+            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
+        Self::MAP(value)
+    }
+}
+
+impl<K> From<HashMap<K, ScalarValue>> for ParamValues
+where
+    K: Into<String>,
+{
+    fn from(value: HashMap<K, ScalarValue>) -> Self {
+        let value: HashMap<String, ScalarValue> =
+            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
         Self::MAP(value)
     }
 }

datafusion/core/src/dataframe/mod.rs

@@ -1215,7 +1215,7 @@ impl DataFrame {
     ///   .with_param_values(vec![
     ///      // value at index 0 --> $1
     ///      ScalarValue::from(2i64)
-    ///    ].into())?
+    ///    ])?
     ///   .collect()
     ///   .await?;
     /// assert_batches_eq!(
@@ -1231,7 +1231,7 @@ impl DataFrame {
     /// # Ok(())
     /// # }
     /// ```
-    pub fn with_param_values(self, query_values: ParamValues) -> Result<Self> {
+    pub fn with_param_values(self, query_values: impl Into<ParamValues>) -> Result<Self> {
         let plan = self.plan.with_param_values(query_values)?;
         Ok(Self::new(self.session_state, plan))
     }

datafusion/core/tests/sql/select.rs

@@ -16,8 +16,7 @@
 // under the License.
 
 use super::*;
-use datafusion_common::{ParamValues, ScalarValue};
-use std::collections::HashMap;
+use datafusion_common::ScalarValue;
 use tempfile::TempDir;
 
 #[tokio::test]
@@ -491,8 +490,7 @@ async fn test_prepare_statement() -> Result<()> {
         ctx.sql("PREPARE my_plan(INT, DOUBLE) AS SELECT c1, c2 FROM test WHERE c1 > $2 AND c1 < $1").await?;
 
     // prepare logical plan to logical plan without parameters
-    let param_values: ParamValues =
-        vec![ScalarValue::Int32(Some(3)), ScalarValue::Float64(Some(0.0))].into();
+    let param_values = vec![ScalarValue::Int32(Some(3)), ScalarValue::Float64(Some(0.0))];
     let dataframe = dataframe.with_param_values(param_values)?;
     let results = dataframe.collect().await?;
 
@@ -535,18 +533,15 @@ async fn test_named_query_parameters() -> Result<()> {
 
     // sql to statement then to logical plan with parameters
     // c1 defined as UINT32, c2 defined as UInt64
-    let dataframe = ctx
+    let results = ctx
         .sql("SELECT c1, c2 FROM test WHERE c1 > $coo AND c1 < $foo")
+        .await?
+        .with_param_values(vec![
+            ("foo", ScalarValue::UInt32(Some(3))),
+            ("coo", ScalarValue::UInt32(Some(0))),
+        ])?
+        .collect()
         .await?;
-
-    // to logical plan without parameters
-    let param_values = vec![
-        ("foo".to_string(), ScalarValue::UInt32(Some(3))),
-        ("coo".to_string(), ScalarValue::UInt32(Some(0))),
-    ];
-    let param_values: HashMap<_, _> = param_values.into_iter().collect();
-    let dataframe = dataframe.with_param_values(param_values.into())?;
-    let results = dataframe.collect().await?;
     let expected = vec![
         "+----+----+",
         "| c1 | c2 |",

datafusion/expr/src/logical_plan/plan.rs

@@ -981,14 +981,18 @@ impl LogicalPlan {
     /// // Fill in the parameter $1 with a literal 3
     /// let plan = plan.with_param_values(vec![
     ///   ScalarValue::from(3i32) // value at index 0 --> $1
-    /// ].into()).unwrap();
+    /// ]).unwrap();
     ///
     /// assert_eq!("Filter: t1.id = Int32(3)\
     ///             \n  TableScan: t1",
     ///             plan.display_indent().to_string()
     ///  );
     /// ```
-    pub fn with_param_values(self, param_values: ParamValues) -> Result<LogicalPlan> {
+    pub fn with_param_values(
+        self,
+        param_values: impl Into<ParamValues>,
+    ) -> Result<LogicalPlan> {
+        let param_values = param_values.into();
         match self {
             LogicalPlan::Prepare(prepare_lp) => {
                 param_values.verify(&prepare_lp.data_types)?;

datafusion/sql/tests/sql_integration.rs

@@ -2678,7 +2678,7 @@ fn prepare_stmt_quick_test(
 
 fn prepare_stmt_replace_params_quick_test(
     plan: LogicalPlan,
-    param_values: ParamValues,
+    param_values: impl Into<ParamValues>,
     expected_plan: &str,
 ) -> LogicalPlan {
     // replace params
@@ -3708,7 +3708,7 @@ fn test_prepare_statement_to_plan_no_param() {
 
     ///////////////////
     // replace params with values
-    let param_values = vec![ScalarValue::Int32(Some(10))].into();
+    let param_values = vec![ScalarValue::Int32(Some(10))];
     let expected_plan = "Projection: person.id, person.age\
         \n  Filter: person.age = Int64(10)\
         \n    TableScan: person";
@@ -3730,7 +3730,7 @@ fn test_prepare_statement_to_plan_no_param() {
 
     ///////////////////
     // replace params with values
-    let param_values = vec![].into();
+    let param_values: Vec<ScalarValue> = vec![];
     let expected_plan = "Projection: person.id, person.age\
         \n  Filter: person.age = Int64(10)\
         \n    TableScan: person";
@@ -3744,7 +3744,7 @@ fn test_prepare_statement_to_plan_one_param_no_value_panic() {
     let sql = "PREPARE my_plan(INT) AS SELECT id, age  FROM person WHERE age = 10";
     let plan = logical_plan(sql).unwrap();
     // declare 1 param but provide 0
-    let param_values = vec![].into();
+    let param_values: Vec<ScalarValue> = vec![];
     assert_eq!(
         plan.with_param_values(param_values)
             .unwrap_err()
@@ -3759,7 +3759,7 @@ fn test_prepare_statement_to_plan_one_param_one_value_different_type_panic() {
     let sql = "PREPARE my_plan(INT) AS SELECT id, age  FROM person WHERE age = 10";
     let plan = logical_plan(sql).unwrap();
     // declare 1 param but provide 0
-    let param_values = vec![ScalarValue::Float64(Some(20.0))].into();
+    let param_values = vec![ScalarValue::Float64(Some(20.0))];
     assert_eq!(
         plan.with_param_values(param_values)
             .unwrap_err()
@@ -3774,7 +3774,7 @@ fn test_prepare_statement_to_plan_no_param_on_value_panic() {
     let sql = "PREPARE my_plan AS SELECT id, age  FROM person WHERE age = 10";
     let plan = logical_plan(sql).unwrap();
     // declare 1 param but provide 0
-    let param_values = vec![ScalarValue::Int32(Some(10))].into();
+    let param_values = vec![ScalarValue::Int32(Some(10))];
     assert_eq!(
         plan.with_param_values(param_values)
             .unwrap_err()
@@ -3795,7 +3795,7 @@ fn test_prepare_statement_to_plan_params_as_constants() {
 
     ///////////////////
     // replace params with values
-    let param_values = vec![ScalarValue::Int32(Some(10))].into();
+    let param_values = vec![ScalarValue::Int32(Some(10))];
     let expected_plan = "Projection: Int32(10)\n  EmptyRelation";
 
     prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan);
@@ -3811,7 +3811,7 @@ fn test_prepare_statement_to_plan_params_as_constants() {
 
     ///////////////////
     // replace params with values
-    let param_values = vec![ScalarValue::Int32(Some(10))].into();
+    let param_values = vec![ScalarValue::Int32(Some(10))];
     let expected_plan = "Projection: Int64(1) + Int32(10)\n  EmptyRelation";
 
     prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan);
@@ -3830,8 +3830,7 @@ fn test_prepare_statement_to_plan_params_as_constants() {
     let param_values = vec![
         ScalarValue::Int32(Some(10)),
         ScalarValue::Float64(Some(10.0)),
-    ]
-    .into();
+    ];
     let expected_plan = "Projection: Int64(1) + Int32(10) + Float64(10)\n  EmptyRelation";
 
     prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan);
@@ -4068,7 +4067,7 @@ fn test_prepare_statement_to_plan_one_param() {
 
     ///////////////////
     // replace params with values
-    let param_values = vec![ScalarValue::Int32(Some(10))].into();
+    let param_values = vec![ScalarValue::Int32(Some(10))];
     let expected_plan = "Projection: person.id, person.age\
         \n  Filter: person.age = Int32(10)\
         \n    TableScan: person";
@@ -4093,7 +4092,7 @@ fn test_prepare_statement_to_plan_data_type() {
 
     ///////////////////
     // replace params with values still succeed and use Float64
-    let param_values = vec![ScalarValue::Float64(Some(10.0))].into();
+    let param_values = vec![ScalarValue::Float64(Some(10.0))];
     let expected_plan = "Projection: person.id, person.age\
         \n  Filter: person.age = Float64(10)\
         \n    TableScan: person";
@@ -4126,8 +4125,7 @@ fn test_prepare_statement_to_plan_multi_params() {
         ScalarValue::Int32(Some(20)),
         ScalarValue::Float64(Some(200.0)),
         ScalarValue::Utf8(Some("xyz".to_string())),
-    ]
-    .into();
+    ];
     let expected_plan =
             "Projection: person.id, person.age, Utf8(\"xyz\")\
         \n  Filter: person.age IN ([Int32(10), Int32(20)]) AND person.salary > Float64(100) AND person.salary < Float64(200) OR person.first_name < Utf8(\"abc\")\
@@ -4164,8 +4162,7 @@ fn test_prepare_statement_to_plan_having() {
         ScalarValue::Float64(Some(100.0)),
         ScalarValue::Float64(Some(200.0)),
         ScalarValue::Float64(Some(300.0)),
-    ]
-    .into();
+    ];
     let expected_plan =
             "Projection: person.id, SUM(person.age)\
         \n  Filter: SUM(person.age) < Int32(10) AND SUM(person.age) > Int64(10) OR SUM(person.age) IN ([Float64(200), Float64(300)])\
@@ -4195,8 +4192,7 @@ fn test_prepare_statement_to_plan_value_list() {
     let param_values = vec![
         ScalarValue::Utf8(Some("a".to_string())),
         ScalarValue::Utf8(Some("b".to_string())),
-    ]
-    .into();
+    ];
     let expected_plan = "Projection: t.num, t.letter\
         \n  SubqueryAlias: t\
         \n    Projection: column1 AS num, column2 AS letter\