datafusion_common/
param_value.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::error::{_plan_datafusion_err, _plan_err};
19use crate::{Result, ScalarValue};
20use arrow::datatypes::DataType;
21use std::collections::HashMap;
22
23/// The parameter value corresponding to the placeholder
24#[derive(Debug, Clone)]
25pub enum ParamValues {
26    /// For positional query parameters, like `SELECT * FROM test WHERE a > $1 AND b = $2`
27    List(Vec<ScalarValue>),
28    /// For named query parameters, like `SELECT * FROM test WHERE a > $foo AND b = $goo`
29    Map(HashMap<String, ScalarValue>),
30}
31
32impl ParamValues {
33    /// Verify parameter list length and type
34    pub fn verify(&self, expect: &[DataType]) -> Result<()> {
35        match self {
36            ParamValues::List(list) => {
37                // Verify if the number of params matches the number of values
38                if expect.len() != list.len() {
39                    return _plan_err!(
40                        "Expected {} parameters, got {}",
41                        expect.len(),
42                        list.len()
43                    );
44                }
45
46                // Verify if the types of the params matches the types of the values
47                let iter = expect.iter().zip(list.iter());
48                for (i, (param_type, value)) in iter.enumerate() {
49                    if *param_type != value.data_type() {
50                        return _plan_err!(
51                            "Expected parameter of type {:?}, got {:?} at index {}",
52                            param_type,
53                            value.data_type(),
54                            i
55                        );
56                    }
57                }
58                Ok(())
59            }
60            ParamValues::Map(_) => {
61                // If it is a named query, variables can be reused,
62                // but the lengths are not necessarily equal
63                Ok(())
64            }
65        }
66    }
67
68    pub fn get_placeholders_with_values(&self, id: &str) -> Result<ScalarValue> {
69        match self {
70            ParamValues::List(list) => {
71                if id.is_empty() {
72                    return _plan_err!("Empty placeholder id");
73                }
74                // convert id (in format $1, $2, ..) to idx (0, 1, ..)
75                let idx = id[1..]
76                    .parse::<usize>()
77                    .map_err(|e| {
78                        _plan_datafusion_err!("Failed to parse placeholder id: {e}")
79                    })?
80                    .checked_sub(1);
81                // value at the idx-th position in param_values should be the value for the placeholder
82                let value = idx.and_then(|idx| list.get(idx)).ok_or_else(|| {
83                    _plan_datafusion_err!("No value found for placeholder with id {id}")
84                })?;
85                Ok(value.clone())
86            }
87            ParamValues::Map(map) => {
88                // convert name (in format $a, $b, ..) to mapped values (a, b, ..)
89                let name = &id[1..];
90                // value at the name position in param_values should be the value for the placeholder
91                let value = map.get(name).ok_or_else(|| {
92                    _plan_datafusion_err!("No value found for placeholder with name {id}")
93                })?;
94                Ok(value.clone())
95            }
96        }
97    }
98}
99
100impl From<Vec<ScalarValue>> for ParamValues {
101    fn from(value: Vec<ScalarValue>) -> Self {
102        Self::List(value)
103    }
104}
105
106impl<K> From<Vec<(K, ScalarValue)>> for ParamValues
107where
108    K: Into<String>,
109{
110    fn from(value: Vec<(K, ScalarValue)>) -> Self {
111        let value: HashMap<String, ScalarValue> =
112            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
113        Self::Map(value)
114    }
115}
116
117impl<K> From<HashMap<K, ScalarValue>> for ParamValues
118where
119    K: Into<String>,
120{
121    fn from(value: HashMap<K, ScalarValue>) -> Self {
122        let value: HashMap<String, ScalarValue> =
123            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
124        Self::Map(value)
125    }
126}
datafusion_common/param_value.rs

datafusion_common/
param_value.rs
