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::metadata::{check_metadata_with_storage_equal, ScalarAndMetadata};
20use crate::{Result, ScalarValue};
21use arrow::datatypes::{DataType, Field, FieldRef};
22use std::collections::HashMap;
23
24/// The parameter value corresponding to the placeholder
25#[derive(Debug, Clone)]
26pub enum ParamValues {
27    /// For positional query parameters, like `SELECT * FROM test WHERE a > $1 AND b = $2`
28    List(Vec<ScalarAndMetadata>),
29    /// For named query parameters, like `SELECT * FROM test WHERE a > $foo AND b = $goo`
30    Map(HashMap<String, ScalarAndMetadata>),
31}
32
33impl ParamValues {
34    /// Verify parameter list length and DataType
35    ///
36    /// Use [`ParamValues::verify_fields`] to ensure field metadata is considered when
37    /// computing type equality.
38    #[deprecated(since = "51.0.0", note = "Use verify_fields instead")]
39    pub fn verify(&self, expect: &[DataType]) -> Result<()> {
40        // make dummy Fields
41        let expect = expect
42            .iter()
43            .map(|dt| Field::new("", dt.clone(), true).into())
44            .collect::<Vec<_>>();
45        self.verify_fields(&expect)
46    }
47
48    /// Verify parameter list length and type
49    pub fn verify_fields(&self, expect: &[FieldRef]) -> Result<()> {
50        match self {
51            ParamValues::List(list) => {
52                // Verify if the number of params matches the number of values
53                if expect.len() != list.len() {
54                    return _plan_err!(
55                        "Expected {} parameters, got {}",
56                        expect.len(),
57                        list.len()
58                    );
59                }
60
61                // Verify if the types of the params matches the types of the values
62                let iter = expect.iter().zip(list.iter());
63                for (i, (param_type, lit)) in iter.enumerate() {
64                    check_metadata_with_storage_equal(
65                        (
66                            &lit.value.data_type(),
67                            lit.metadata.as_ref().map(|m| m.to_hashmap()).as_ref(),
68                        ),
69                        (param_type.data_type(), Some(param_type.metadata())),
70                        "parameter",
71                        &format!(" at index {i}"),
72                    )?;
73                }
74                Ok(())
75            }
76            ParamValues::Map(_) => {
77                // If it is a named query, variables can be reused,
78                // but the lengths are not necessarily equal
79                Ok(())
80            }
81        }
82    }
83
84    pub fn get_placeholders_with_values(&self, id: &str) -> Result<ScalarAndMetadata> {
85        match self {
86            ParamValues::List(list) => {
87                if id.is_empty() {
88                    return _plan_err!("Empty placeholder id");
89                }
90                // convert id (in format $1, $2, ..) to idx (0, 1, ..)
91                let idx = id[1..]
92                    .parse::<usize>()
93                    .map_err(|e| {
94                        _plan_datafusion_err!("Failed to parse placeholder id: {e}")
95                    })?
96                    .checked_sub(1);
97                // value at the idx-th position in param_values should be the value for the placeholder
98                let value = idx.and_then(|idx| list.get(idx)).ok_or_else(|| {
99                    _plan_datafusion_err!("No value found for placeholder with id {id}")
100                })?;
101                Ok(value.clone())
102            }
103            ParamValues::Map(map) => {
104                // convert name (in format $a, $b, ..) to mapped values (a, b, ..)
105                let name = &id[1..];
106                // value at the name position in param_values should be the value for the placeholder
107                let value = map.get(name).ok_or_else(|| {
108                    _plan_datafusion_err!("No value found for placeholder with name {id}")
109                })?;
110                Ok(value.clone())
111            }
112        }
113    }
114}
115
116impl From<Vec<ScalarValue>> for ParamValues {
117    fn from(value: Vec<ScalarValue>) -> Self {
118        Self::List(value.into_iter().map(ScalarAndMetadata::from).collect())
119    }
120}
121
122impl<K> From<Vec<(K, ScalarValue)>> for ParamValues
123where
124    K: Into<String>,
125{
126    fn from(value: Vec<(K, ScalarValue)>) -> Self {
127        let value: HashMap<String, ScalarAndMetadata> = value
128            .into_iter()
129            .map(|(k, v)| (k.into(), ScalarAndMetadata::from(v)))
130            .collect();
131        Self::Map(value)
132    }
133}
134
135impl<K> From<HashMap<K, ScalarValue>> for ParamValues
136where
137    K: Into<String>,
138{
139    fn from(value: HashMap<K, ScalarValue>) -> Self {
140        let value: HashMap<String, ScalarAndMetadata> = value
141            .into_iter()
142            .map(|(k, v)| (k.into(), ScalarAndMetadata::from(v)))
143            .collect();
144        Self::Map(value)
145    }
146}