You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by al...@apache.org on 2025/05/25 11:18:37 UTC
(datafusion) branch main updated: chore: Reduce repetition in the parameter type inference tests (#16079)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 3b551e93a9 chore: Reduce repetition in the parameter type inference tests (#16079)
3b551e93a9 is described below
commit 3b551e93a9e87b81d2ee7dd14477f30383bb052d
Author: jsai28 <54...@users.noreply.github.com>
AuthorDate: Sun May 25 05:18:33 2025 -0600
chore: Reduce repetition in the parameter type inference tests (#16079)
* added test
* added parameterTest
* cargo fmt
* Update sql_integration.rs
* allow needless_lifetimes
* remove needless lifetime
* update some tests
* move to params.rs
---
datafusion/sql/tests/cases/params.rs | 412 ++++++++++++++---------------------
1 file changed, 167 insertions(+), 245 deletions(-)
diff --git a/datafusion/sql/tests/cases/params.rs b/datafusion/sql/tests/cases/params.rs
index 7a4fbba147..b3cc49c310 100644
--- a/datafusion/sql/tests/cases/params.rs
+++ b/datafusion/sql/tests/cases/params.rs
@@ -22,6 +22,34 @@ use datafusion_expr::{LogicalPlan, Prepare, Statement};
use insta::assert_snapshot;
use std::collections::HashMap;
+pub struct ParameterTest<'a> {
+ pub sql: &'a str,
+ pub expected_types: Vec<(&'a str, Option<DataType>)>,
+ pub param_values: Vec<ScalarValue>,
+}
+
+impl ParameterTest<'_> {
+ pub fn run(&self) -> String {
+ let plan = logical_plan(self.sql).unwrap();
+
+ let actual_types = plan.get_parameter_types().unwrap();
+ let expected_types: HashMap<String, Option<DataType>> = self
+ .expected_types
+ .iter()
+ .map(|(k, v)| (k.to_string(), v.clone()))
+ .collect();
+
+ assert_eq!(actual_types, expected_types);
+
+ let plan_with_params = plan
+ .clone()
+ .with_param_values(self.param_values.clone())
+ .unwrap();
+
+ format!("** Initial Plan:\n{plan}\n** Final Plan:\n{plan_with_params}")
+ }
+}
+
fn generate_prepare_stmt_and_data_types(sql: &str) -> (LogicalPlan, String) {
let plan = logical_plan(sql).unwrap();
let data_types = match &plan {
@@ -311,31 +339,22 @@ fn test_prepare_statement_to_plan_params_as_constants() {
#[test]
fn test_infer_types_from_join() {
- let sql =
- "SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1";
+ let test = ParameterTest {
+ sql:
+ "SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1",
+ expected_types: vec![("$1", Some(DataType::Int32))],
+ param_values: vec![ScalarValue::Int32(Some(10))],
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
- @r#"
+ test.run(),
+ @r"
+ ** Initial Plan:
Projection: person.id, orders.order_id
Inner Join: Filter: person.id = orders.customer_id AND person.age = $1
TableScan: person
TableScan: orders
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, orders.order_id
Inner Join: Filter: person.id = orders.customer_id AND person.age = Int32(10)
TableScan: person
@@ -346,64 +365,46 @@ fn test_infer_types_from_join() {
#[test]
fn test_prepare_statement_infer_types_from_join() {
- let sql =
- "PREPARE my_plan AS SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1";
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1",
+ expected_types: vec![("$1", Some(DataType::Int32))],
+ param_values: vec![ScalarValue::Int32(Some(10))]
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [Int32]
Projection: person.id, orders.order_id
Inner Join: Filter: person.id = orders.customer_id AND person.age = $1
TableScan: person
TableScan: orders
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, orders.order_id
Inner Join: Filter: person.id = orders.customer_id AND person.age = Int32(10)
TableScan: person
TableScan: orders
- "
+ "#
);
}
#[test]
fn test_infer_types_from_predicate() {
- let sql = "SELECT id, age FROM person WHERE age = $1";
- let plan = logical_plan(sql).unwrap();
+ let test = ParameterTest {
+ sql: "SELECT id, age FROM person WHERE age = $1",
+ expected_types: vec![("$1", Some(DataType::Int32))],
+ param_values: vec![ScalarValue::Int32(Some(10))],
+ };
+
assert_snapshot!(
- plan,
- @r#"
+ test.run(),
+ @r"
+ ** Initial Plan:
Projection: person.id, person.age
Filter: person.age = $1
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age = Int32(10)
TableScan: person
@@ -413,64 +414,46 @@ fn test_infer_types_from_predicate() {
#[test]
fn test_prepare_statement_infer_types_from_predicate() {
- let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age = $1";
- let plan = logical_plan(sql).unwrap();
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS SELECT id, age FROM person WHERE age = $1",
+ expected_types: vec![("$1", Some(DataType::Int32))],
+ param_values: vec![ScalarValue::Int32(Some(10))],
+ };
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [Int32]
Projection: person.id, person.age
Filter: person.age = $1
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age = Int32(10)
TableScan: person
- "
+ "#
);
}
#[test]
fn test_infer_types_from_between_predicate() {
- let sql = "SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
+ let test = ParameterTest {
+ sql: "SELECT id, age FROM person WHERE age BETWEEN $1 AND $2",
+ expected_types: vec![
+ ("$1", Some(DataType::Int32)),
+ ("$2", Some(DataType::Int32)),
+ ],
+ param_values: vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(30))],
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
- @r#"
+ test.run(),
+ @r"
+ ** Initial Plan:
Projection: person.id, person.age
Filter: person.age BETWEEN $1 AND $2
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::Int32)),
- ("$2".to_string(), Some(DataType::Int32)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(30))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age BETWEEN Int32(10) AND Int32(30)
TableScan: person
@@ -480,48 +463,42 @@ fn test_infer_types_from_between_predicate() {
#[test]
fn test_prepare_statement_infer_types_from_between_predicate() {
- let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
-
- let plan = logical_plan(sql).unwrap();
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS SELECT id, age FROM person WHERE age BETWEEN $1 AND $2",
+ expected_types: vec![
+ ("$1", Some(DataType::Int32)),
+ ("$2", Some(DataType::Int32)),
+ ],
+ param_values: vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(30))],
+ };
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [Int32, Int32]
Projection: person.id, person.age
Filter: person.age BETWEEN $1 AND $2
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::Int32)),
- ("$2".to_string(), Some(DataType::Int32)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(30))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age BETWEEN Int32(10) AND Int32(30)
TableScan: person
- "
+ "#
);
}
#[test]
fn test_infer_types_subquery() {
- let sql = "SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)";
+ let test = ParameterTest {
+ sql: "SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)",
+ expected_types: vec![("$1", Some(DataType::UInt32))],
+ param_values: vec![ScalarValue::UInt32(Some(10))]
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
- @r#"
+ test.run(),
+ @r"
+ ** Initial Plan:
Projection: person.id, person.age
Filter: person.age = (<subquery>)
Subquery:
@@ -530,20 +507,7 @@ fn test_infer_types_subquery() {
Filter: person.id = $1
TableScan: person
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::UInt32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::UInt32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age = (<subquery>)
Subquery:
@@ -552,18 +516,22 @@ fn test_infer_types_subquery() {
Filter: person.id = UInt32(10)
TableScan: person
TableScan: person
- "
+ "
);
}
#[test]
fn test_prepare_statement_infer_types_subquery() {
- let sql = "PREPARE my_plan AS SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)";
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)",
+ expected_types: vec![("$1", Some(DataType::UInt32))],
+ param_values: vec![ScalarValue::UInt32(Some(10))]
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [UInt32]
Projection: person.id, person.age
Filter: person.age = (<subquery>)
@@ -573,20 +541,7 @@ fn test_prepare_statement_infer_types_subquery() {
Filter: person.id = $1
TableScan: person
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([("$1".to_string(), Some(DataType::UInt32))]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::UInt32(Some(10))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Projection: person.id, person.age
Filter: person.age = (<subquery>)
Subquery:
@@ -595,116 +550,91 @@ fn test_prepare_statement_infer_types_subquery() {
Filter: person.id = UInt32(10)
TableScan: person
TableScan: person
- "
+ "#
);
}
#[test]
fn test_update_infer() {
- let sql = "update person set age=$1 where id=$2";
+ let test = ParameterTest {
+ sql: "update person set age=$1 where id=$2",
+ expected_types: vec![
+ ("$1", Some(DataType::Int32)),
+ ("$2", Some(DataType::UInt32)),
+ ],
+ param_values: vec![ScalarValue::Int32(Some(42)), ScalarValue::UInt32(Some(1))],
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
- @r#"
+ test.run(),
+ @r"
+ ** Initial Plan:
Dml: op=[Update] table=[person]
Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, $1 AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
Filter: person.id = $2
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::Int32)),
- ("$2".to_string(), Some(DataType::UInt32)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(42)), ScalarValue::UInt32(Some(1))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Dml: op=[Update] table=[person]
Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, Int32(42) AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
Filter: person.id = UInt32(1)
TableScan: person
- "
+ "
);
}
#[test]
fn test_prepare_statement_update_infer() {
- let sql = "PREPARE my_plan AS update person set age=$1 where id=$2";
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS update person set age=$1 where id=$2",
+ expected_types: vec![
+ ("$1", Some(DataType::Int32)),
+ ("$2", Some(DataType::UInt32)),
+ ],
+ param_values: vec![ScalarValue::Int32(Some(42)), ScalarValue::UInt32(Some(1))],
+ };
- let plan = logical_plan(sql).unwrap();
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [Int32, UInt32]
Dml: op=[Update] table=[person]
Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, $1 AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
Filter: person.id = $2
TableScan: person
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::Int32)),
- ("$2".to_string(), Some(DataType::UInt32)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![ScalarValue::Int32(Some(42)), ScalarValue::UInt32(Some(1))];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
-
- assert_snapshot!(
- plan_with_params,
- @r"
+ ** Final Plan:
Dml: op=[Update] table=[person]
Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, Int32(42) AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
Filter: person.id = UInt32(1)
TableScan: person
- "
+ "#
);
}
#[test]
fn test_insert_infer() {
- let sql = "insert into person (id, first_name, last_name) values ($1, $2, $3)";
- let plan = logical_plan(sql).unwrap();
+ let test = ParameterTest {
+ sql: "insert into person (id, first_name, last_name) values ($1, $2, $3)",
+ expected_types: vec![
+ ("$1", Some(DataType::UInt32)),
+ ("$2", Some(DataType::Utf8)),
+ ("$3", Some(DataType::Utf8)),
+ ],
+ param_values: vec![
+ ScalarValue::UInt32(Some(1)),
+ ScalarValue::from("Alan"),
+ ScalarValue::from("Turing"),
+ ],
+ };
+
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Dml: op=[Insert Into] table=[person]
Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
Values: ($1, $2, $3)
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::UInt32)),
- ("$2".to_string(), Some(DataType::Utf8)),
- ("$3".to_string(), Some(DataType::Utf8)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![
- ScalarValue::UInt32(Some(1)),
- ScalarValue::from("Alan"),
- ScalarValue::from("Turing"),
- ];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
- assert_snapshot!(
- plan_with_params,
- @r#"
+ ** Final Plan:
Dml: op=[Insert Into] table=[person]
Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
Values: (UInt32(1) AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
@@ -714,36 +644,28 @@ fn test_insert_infer() {
#[test]
fn test_prepare_statement_insert_infer() {
- let sql = "PREPARE my_plan AS insert into person (id, first_name, last_name) values ($1, $2, $3)";
- let plan = logical_plan(sql).unwrap();
+ let test = ParameterTest {
+ sql: "PREPARE my_plan AS insert into person (id, first_name, last_name) values ($1, $2, $3)",
+ expected_types: vec![
+ ("$1", Some(DataType::UInt32)),
+ ("$2", Some(DataType::Utf8)),
+ ("$3", Some(DataType::Utf8)),
+ ],
+ param_values: vec![
+ ScalarValue::UInt32(Some(1)),
+ ScalarValue::from("Alan"),
+ ScalarValue::from("Turing"),
+ ]
+ };
assert_snapshot!(
- plan,
+ test.run(),
@r#"
+ ** Initial Plan:
Prepare: "my_plan" [UInt32, Utf8, Utf8]
Dml: op=[Insert Into] table=[person]
Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
Values: ($1, $2, $3)
- "#
- );
-
- let actual_types = plan.get_parameter_types().unwrap();
- let expected_types = HashMap::from([
- ("$1".to_string(), Some(DataType::UInt32)),
- ("$2".to_string(), Some(DataType::Utf8)),
- ("$3".to_string(), Some(DataType::Utf8)),
- ]);
- assert_eq!(actual_types, expected_types);
-
- // replace params with values
- let param_values = vec![
- ScalarValue::UInt32(Some(1)),
- ScalarValue::from("Alan"),
- ScalarValue::from("Turing"),
- ];
- let plan_with_params = plan.with_param_values(param_values).unwrap();
- assert_snapshot!(
- plan_with_params,
- @r#"
+ ** Final Plan:
Dml: op=[Insert Into] table=[person]
Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
Values: (UInt32(1) AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org