fix test

Author: paleolimbot

datafusion/core/tests/sql/select.rs

@@ -18,6 +18,7 @@
 use std::collections::HashMap;
 
 use super::*;
+use datafusion::assert_batches_eq;
 use datafusion_common::{metadata::Literal, ParamValues, ScalarValue};
 use insta::assert_snapshot;
 
@@ -321,35 +322,44 @@ async fn test_named_parameter_not_bound() -> Result<()> {
 
 #[tokio::test]
 async fn test_query_parameters_with_metadata() -> Result<()> {
-    let tmp_dir = TempDir::new()?;
-    let partition_count = 4;
-    let ctx = create_ctx_with_partition(&tmp_dir, partition_count).await?;
+    let ctx = SessionContext::new();
 
-    let metadata0 = HashMap::from([(
-        "some_key".to_string(),
-        "some_value".to_string(),
-    )]);
-    let metadata1 = HashMap::from([(
-        "some_other_key".to_string(),
-        "some_other_value".to_string(),
-    )]);
+    let df = ctx.sql("SELECT $1, $2").await.unwrap();
 
-    // sql to statement then to logical plan with parameters
-    let df = ctx.sql("SELECT $1, $2").await?;
+    let metadata1 = HashMap::from([("some_key".to_string(), "some_value".to_string())]);
+    let metadata2 =
+        HashMap::from([("some_other_key".to_string(), "some_other_value".to_string())]);
+
+    let df_with_params_replaced = df
+        .with_param_values(ParamValues::List(vec![
+            Literal::new(ScalarValue::UInt32(Some(1)), Some(metadata1.clone().into())),
+            Literal::new(
+                ScalarValue::Utf8(Some("two".to_string())),
+                Some(metadata2.clone().into()),
+            ),
+        ]))
+        .unwrap();
 
-    let df_with_params_replaced = df.with_param_values(ParamValues::List(vec![
-        Literal::new(ScalarValue::UInt32(Some(3)), Some(metadata0.clone().into())),
-        Literal::new(
-            ScalarValue::Utf8(Some("bar_value".to_string())),
-            Some(metadata1.clone().into()),
-        ),
-    ]))?;
+    // df_with_params_replaced.schema() is not correct here
+    // https://github.com/apache/datafusion/issues/18102
+    let batches = df_with_params_replaced.clone().collect().await.unwrap();
+    let schema = batches[0].schema();
 
-    let schema = df_with_params_replaced.schema();
     assert_eq!(schema.field(0).data_type(), &DataType::UInt32);
-    assert_eq!(schema.field(0).metadata(), &metadata0);
+    assert_eq!(schema.field(0).metadata(), &metadata1);
     assert_eq!(schema.field(1).data_type(), &DataType::Utf8);
-    assert_eq!(schema.field(1).metadata(), &metadata1);
+    assert_eq!(schema.field(1).metadata(), &metadata2);
+
+    assert_batches_eq!(
+        vec![
+            "+----+-----+",
+            "| $1 | $2  |",
+            "+----+-----+",
+            "| 1  | two |",
+            "+----+-----+",
+        ],
+        &batches
+    );
 
     Ok(())
 }