-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support named query parameters #8384
Changes from 28 commits
7afeb8b
6332bec
cc5e0c7
a114310
928c811
839093e
a836cde
5648dc7
a670409
22894a3
73a59d2
46409c2
8a86a4c
cf5c584
62ae9b9
da02fa2
d98eb2e
79e7216
ba51abd
2468f52
180c303
68980ba
9411940
ba28346
df0942f
d6381c2
69a2080
49e617e
5b8c3b7
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,149 @@ | ||
// Licensed to the Apache Software Foundation (ASF) under one | ||
// or more contributor license agreements. See the NOTICE file | ||
// distributed with this work for additional information | ||
// regarding copyright ownership. The ASF licenses this file | ||
// to you under the Apache License, Version 2.0 (the | ||
// "License"); you may not use this file except in compliance | ||
// with the License. You may obtain a copy of the License at | ||
// | ||
// http://www.apache.org/licenses/LICENSE-2.0 | ||
// | ||
// Unless required by applicable law or agreed to in writing, | ||
// software distributed under the License is distributed on an | ||
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
// KIND, either express or implied. See the License for the | ||
// specific language governing permissions and limitations | ||
// under the License. | ||
|
||
use crate::error::{_internal_err, _plan_err}; | ||
use crate::{DataFusionError, Result, ScalarValue}; | ||
use arrow_schema::DataType; | ||
use std::collections::HashMap; | ||
|
||
/// The parameter value corresponding to the placeholder | ||
#[derive(Debug, Clone)] | ||
pub enum ParamValues { | ||
/// for positional query parameters, like select * from test where a > $1 and b = $2 | ||
LIST(Vec<ScalarValue>), | ||
/// for named query parameters, like select * from test where a > $foo and b = $goo | ||
MAP(HashMap<String, ScalarValue>), | ||
} | ||
|
||
impl ParamValues { | ||
/// Verify parameter list length and type | ||
pub fn verify(&self, expect: &Vec<DataType>) -> Result<()> { | ||
match self { | ||
ParamValues::LIST(list) => { | ||
// Verify if the number of params matches the number of values | ||
if expect.len() != list.len() { | ||
return _plan_err!( | ||
"Expected {} parameters, got {}", | ||
expect.len(), | ||
list.len() | ||
); | ||
} | ||
|
||
// Verify if the types of the params matches the types of the values | ||
let iter = expect.iter().zip(list.iter()); | ||
for (i, (param_type, value)) in iter.enumerate() { | ||
if *param_type != value.data_type() { | ||
return _plan_err!( | ||
"Expected parameter of type {:?}, got {:?} at index {}", | ||
param_type, | ||
value.data_type(), | ||
i | ||
); | ||
} | ||
} | ||
Ok(()) | ||
} | ||
ParamValues::MAP(_) => { | ||
// If it is a named query, variables can be reused, | ||
// but the lengths are not necessarily equal | ||
Ok(()) | ||
} | ||
} | ||
} | ||
|
||
pub fn get_placeholders_with_values( | ||
&self, | ||
id: &String, | ||
data_type: &Option<DataType>, | ||
) -> Result<ScalarValue> { | ||
match self { | ||
ParamValues::LIST(list) => { | ||
if id.is_empty() || id == "$0" { | ||
return _plan_err!("Empty placeholder id"); | ||
} | ||
// convert id (in format $1, $2, ..) to idx (0, 1, ..) | ||
let idx = id[1..].parse::<usize>().map_err(|e| { | ||
DataFusionError::Internal(format!( | ||
"Failed to parse placeholder id: {e}" | ||
)) | ||
})? - 1; | ||
// value at the idx-th position in param_values should be the value for the placeholder | ||
let value = list.get(idx).ok_or_else(|| { | ||
DataFusionError::Internal(format!( | ||
"No value found for placeholder with id {id}" | ||
)) | ||
})?; | ||
// check if the data type of the value matches the data type of the placeholder | ||
if Some(value.data_type()) != *data_type { | ||
return _internal_err!( | ||
"Placeholder value type mismatch: expected {:?}, got {:?}", | ||
data_type, | ||
value.data_type() | ||
); | ||
} | ||
Ok(value.clone()) | ||
} | ||
ParamValues::MAP(map) => { | ||
// convert name (in format $a, $b, ..) to mapped values (a, b, ..) | ||
let name = &id[1..]; | ||
// value at the name position in param_values should be the value for the placeholder | ||
let value = map.get(name).ok_or_else(|| { | ||
DataFusionError::Internal(format!( | ||
"No value found for placeholder with name {id}" | ||
)) | ||
})?; | ||
// check if the data type of the value matches the data type of the placeholder | ||
if Some(value.data_type()) != *data_type { | ||
return _internal_err!( | ||
"Placeholder value type mismatch: expected {:?}, got {:?}", | ||
data_type, | ||
value.data_type() | ||
); | ||
} | ||
Ok(value.clone()) | ||
} | ||
} | ||
} | ||
} | ||
|
||
impl From<Vec<ScalarValue>> for ParamValues { | ||
fn from(value: Vec<ScalarValue>) -> Self { | ||
Self::LIST(value) | ||
} | ||
} | ||
|
||
impl<K> From<Vec<(K, ScalarValue)>> for ParamValues | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 👌 very nice |
||
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) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -48,7 +48,7 @@ use datafusion_common::tree_node::{ | |
use datafusion_common::{ | ||
aggregate_functional_dependencies, internal_err, plan_err, Column, Constraints, | ||
DFField, DFSchema, DFSchemaRef, DataFusionError, FunctionalDependencies, | ||
OwnedTableReference, Result, ScalarValue, UnnestOptions, | ||
OwnedTableReference, ParamValues, Result, UnnestOptions, | ||
}; | ||
// backwards compatibility | ||
pub use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan}; | ||
|
@@ -990,32 +990,12 @@ impl LogicalPlan { | |
/// ``` | ||
Asura7969 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
pub fn with_param_values( | ||
self, | ||
param_values: Vec<ScalarValue>, | ||
param_values: impl Into<ParamValues>, | ||
) -> Result<LogicalPlan> { | ||
let param_values = param_values.into(); | ||
match self { | ||
LogicalPlan::Prepare(prepare_lp) => { | ||
// Verify if the number of params matches the number of values | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. move to ParamValues( |
||
if prepare_lp.data_types.len() != param_values.len() { | ||
return plan_err!( | ||
"Expected {} parameters, got {}", | ||
prepare_lp.data_types.len(), | ||
param_values.len() | ||
); | ||
} | ||
|
||
// Verify if the types of the params matches the types of the values | ||
let iter = prepare_lp.data_types.iter().zip(param_values.iter()); | ||
for (i, (param_type, value)) in iter.enumerate() { | ||
if *param_type != value.data_type() { | ||
return plan_err!( | ||
"Expected parameter of type {:?}, got {:?} at index {}", | ||
param_type, | ||
value.data_type(), | ||
i | ||
); | ||
} | ||
} | ||
|
||
param_values.verify(&prepare_lp.data_types)?; | ||
let input_plan = prepare_lp.input; | ||
input_plan.replace_params_with_values(¶m_values) | ||
} | ||
|
@@ -1179,7 +1159,7 @@ impl LogicalPlan { | |
/// See [`Self::with_param_values`] for examples and usage | ||
pub fn replace_params_with_values( | ||
&self, | ||
param_values: &[ScalarValue], | ||
param_values: &ParamValues, | ||
) -> Result<LogicalPlan> { | ||
let new_exprs = self | ||
.expressions() | ||
|
@@ -1236,36 +1216,15 @@ impl LogicalPlan { | |
/// corresponding values provided in the params_values | ||
fn replace_placeholders_with_values( | ||
expr: Expr, | ||
param_values: &[ScalarValue], | ||
param_values: &ParamValues, | ||
) -> Result<Expr> { | ||
expr.transform(&|expr| { | ||
match &expr { | ||
Expr::Placeholder(Placeholder { id, data_type }) => { | ||
if id.is_empty() || id == "$0" { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. move to ParamValues( |
||
return plan_err!("Empty placeholder id"); | ||
} | ||
// convert id (in format $1, $2, ..) to idx (0, 1, ..) | ||
let idx = id[1..].parse::<usize>().map_err(|e| { | ||
DataFusionError::Internal(format!( | ||
"Failed to parse placeholder id: {e}" | ||
)) | ||
})? - 1; | ||
// value at the idx-th position in param_values should be the value for the placeholder | ||
let value = param_values.get(idx).ok_or_else(|| { | ||
DataFusionError::Internal(format!( | ||
"No value found for placeholder with id {id}" | ||
)) | ||
})?; | ||
// check if the data type of the value matches the data type of the placeholder | ||
if Some(value.data_type()) != *data_type { | ||
return internal_err!( | ||
"Placeholder value type mismatch: expected {:?}, got {:?}", | ||
data_type, | ||
value.data_type() | ||
); | ||
} | ||
let value = | ||
param_values.get_placeholders_with_values(id, data_type)?; | ||
// Replace the placeholder with the value | ||
Ok(Transformed::Yes(Expr::Literal(value.clone()))) | ||
Ok(Transformed::Yes(Expr::Literal(value))) | ||
} | ||
Expr::ScalarSubquery(qry) => { | ||
let subquery = | ||
|
@@ -2556,7 +2515,7 @@ mod tests { | |
use crate::{col, count, exists, in_subquery, lit, placeholder, GroupingSet}; | ||
use arrow::datatypes::{DataType, Field, Schema}; | ||
use datafusion_common::tree_node::TreeNodeVisitor; | ||
use datafusion_common::{not_impl_err, DFSchema, TableReference}; | ||
use datafusion_common::{not_impl_err, DFSchema, ScalarValue, TableReference}; | ||
use std::collections::HashMap; | ||
|
||
fn employee_schema() -> Schema { | ||
|
@@ -3004,7 +2963,8 @@ digraph { | |
.build() | ||
.unwrap(); | ||
|
||
plan.replace_params_with_values(&[42i32.into()]) | ||
let param_values = vec![ScalarValue::Int32(Some(42))]; | ||
plan.replace_params_with_values(¶m_values.clone().into()) | ||
.expect_err("unexpectedly succeeded to replace an invalid placeholder"); | ||
|
||
// test $0 placeholder | ||
|
@@ -3017,7 +2977,7 @@ digraph { | |
.build() | ||
.unwrap(); | ||
|
||
plan.replace_params_with_values(&[42i32.into()]) | ||
plan.replace_params_with_values(¶m_values.into()) | ||
.expect_err("unexpectedly succeeded to replace an invalid placeholder"); | ||
} | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think it is more appropriate to encapsulate it into a struct