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 51452d6e7d Fix Infer prepare statement type tests (#15743)
51452d6e7d is described below
commit 51452d6e7d9c95ce60aeb9a6834dbe82cdd29a99
Author: Brayan Jules <[email protected]>
AuthorDate: Thu May 15 08:35:12 2025 -0400
Fix Infer prepare statement type tests (#15743)
* draft commit to rolledback changes on function naming and include prepare
clause on the infer types tests
* include data types in plan when it is not included in the prepare
statement
* fix: prepare statement error
* Update datafusion/sql/src/statement.rs
Co-authored-by: Andrew Lamb <[email protected]>
* remove infer types from prepare statement
the infer data type changes in statement will be introduced in a new PR
* fix to show correct output message
* remove white space
* Restore the original tests too
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/sql/tests/sql_integration.rs | 136 ++++++++++++++++++++++++++++++++
1 file changed, 136 insertions(+)
diff --git a/datafusion/sql/tests/sql_integration.rs
b/datafusion/sql/tests/sql_integration.rs
index 2804a1de06..6d08dd4fe5 100644
--- a/datafusion/sql/tests/sql_integration.rs
+++ b/datafusion/sql/tests/sql_integration.rs
@@ -4641,6 +4641,28 @@ 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 plan = logical_plan(sql).unwrap();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_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);
+}
+
#[test]
fn test_infer_types_from_predicate() {
let sql = "SELECT id, age FROM person WHERE age = $1";
@@ -4672,6 +4694,25 @@ 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();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_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);
+}
+
#[test]
fn test_infer_types_from_between_predicate() {
let sql = "SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
@@ -4707,6 +4748,29 @@ 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();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_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);
+}
+
#[test]
fn test_infer_types_subquery() {
let sql = "SELECT id, age FROM person WHERE age = (select max(age) from
person where id = $1)";
@@ -4749,6 +4813,31 @@ fn test_infer_types_subquery() {
);
}
+#[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 plan = logical_plan(sql).unwrap();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_plan" []
+ Projection: person.id, person.age
+ Filter: person.age = (<subquery>)
+ Subquery:
+ Projection: max(person.age)
+ Aggregate: groupBy=[[]], aggr=[[max(person.age)]]
+ 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);
+}
+
#[test]
fn test_update_infer() {
let sql = "update person set age=$1 where id=$2";
@@ -4786,6 +4875,30 @@ fn test_update_infer() {
);
}
+#[test]
+fn test_prepare_statement_update_infer() {
+ let sql = "PREPARE my_plan AS update person set age=$1 where id=$2";
+
+ let plan = logical_plan(sql).unwrap();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_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);
+}
+
#[test]
fn test_insert_infer() {
let sql = "insert into person (id, first_name, last_name) values ($1, $2,
$3)";
@@ -4824,6 +4937,29 @@ 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();
+ assert_snapshot!(
+ plan,
+ @r#"
+ Prepare: "my_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);
+}
+
#[test]
fn test_prepare_statement_to_plan_one_param() {
let sql = "PREPARE my_plan(INT) AS SELECT id, age FROM person WHERE age =
$1";
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
