/Users/andrewlamb/Software/datafusion/datafusion/common/src/param_value.rs
Line | Count | Source (jump to first uncovered line) |
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 | | |
18 | | use crate::error::{_plan_datafusion_err, _plan_err}; |
19 | | use crate::{Result, ScalarValue}; |
20 | | use arrow_schema::DataType; |
21 | | use std::collections::HashMap; |
22 | | |
23 | | /// The parameter value corresponding to the placeholder |
24 | | #[derive(Debug, Clone)] |
25 | | pub 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 | | |
32 | | impl ParamValues { |
33 | | /// Verify parameter list length and type |
34 | 0 | pub fn verify(&self, expect: &[DataType]) -> Result<()> { |
35 | 0 | match self { |
36 | 0 | ParamValues::List(list) => { |
37 | 0 | // Verify if the number of params matches the number of values |
38 | 0 | if expect.len() != list.len() { |
39 | 0 | return _plan_err!( |
40 | 0 | "Expected {} parameters, got {}", |
41 | 0 | expect.len(), |
42 | 0 | list.len() |
43 | 0 | ); |
44 | 0 | } |
45 | 0 |
|
46 | 0 | // Verify if the types of the params matches the types of the values |
47 | 0 | let iter = expect.iter().zip(list.iter()); |
48 | 0 | for (i, (param_type, value)) in iter.enumerate() { |
49 | 0 | if *param_type != value.data_type() { |
50 | 0 | return _plan_err!( |
51 | 0 | "Expected parameter of type {:?}, got {:?} at index {}", |
52 | 0 | param_type, |
53 | 0 | value.data_type(), |
54 | 0 | i |
55 | 0 | ); |
56 | 0 | } |
57 | | } |
58 | 0 | 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 | 0 | Ok(()) |
64 | | } |
65 | | } |
66 | 0 | } |
67 | | |
68 | 0 | pub fn get_placeholders_with_values(&self, id: &str) -> Result<ScalarValue> { |
69 | 0 | match self { |
70 | 0 | ParamValues::List(list) => { |
71 | 0 | if id.is_empty() { |
72 | 0 | return _plan_err!("Empty placeholder id"); |
73 | 0 | } |
74 | | // convert id (in format $1, $2, ..) to idx (0, 1, ..) |
75 | 0 | let idx = id[1..] |
76 | 0 | .parse::<usize>() |
77 | 0 | .map_err(|e| { |
78 | 0 | _plan_datafusion_err!("Failed to parse placeholder id: {e}") |
79 | 0 | })? |
80 | 0 | .checked_sub(1); |
81 | | // value at the idx-th position in param_values should be the value for the placeholder |
82 | 0 | let value = idx.and_then(|idx| list.get(idx)).ok_or_else(|| { |
83 | 0 | _plan_datafusion_err!("No value found for placeholder with id {id}") |
84 | 0 | })?; |
85 | 0 | Ok(value.clone()) |
86 | | } |
87 | 0 | ParamValues::Map(map) => { |
88 | 0 | // convert name (in format $a, $b, ..) to mapped values (a, b, ..) |
89 | 0 | let name = &id[1..]; |
90 | | // value at the name position in param_values should be the value for the placeholder |
91 | 0 | let value = map.get(name).ok_or_else(|| { |
92 | 0 | _plan_datafusion_err!("No value found for placeholder with name {id}") |
93 | 0 | })?; |
94 | 0 | Ok(value.clone()) |
95 | | } |
96 | | } |
97 | 0 | } |
98 | | } |
99 | | |
100 | | impl From<Vec<ScalarValue>> for ParamValues { |
101 | 0 | fn from(value: Vec<ScalarValue>) -> Self { |
102 | 0 | Self::List(value) |
103 | 0 | } |
104 | | } |
105 | | |
106 | | impl<K> From<Vec<(K, ScalarValue)>> for ParamValues |
107 | | where |
108 | | K: Into<String>, |
109 | | { |
110 | 0 | fn from(value: Vec<(K, ScalarValue)>) -> Self { |
111 | 0 | let value: HashMap<String, ScalarValue> = |
112 | 0 | value.into_iter().map(|(k, v)| (k.into(), v)).collect(); |
113 | 0 | Self::Map(value) |
114 | 0 | } |
115 | | } |
116 | | |
117 | | impl<K> From<HashMap<K, ScalarValue>> for ParamValues |
118 | | where |
119 | | K: Into<String>, |
120 | | { |
121 | 0 | fn from(value: HashMap<K, ScalarValue>) -> Self { |
122 | 0 | let value: HashMap<String, ScalarValue> = |
123 | 0 | value.into_iter().map(|(k, v)| (k.into(), v)).collect(); |
124 | 0 | Self::Map(value) |
125 | 0 | } |
126 | | } |