fix cast

Author: friendlymatthew

Cargo.toml

@@ -35,3 +35,6 @@ pedantic = { level = "deny", priority = -1 }
 [[bench]]
 name = "main"
 harness = false
+
+[patch.crates-io]
+datafusion = { git = "https://github.com/pydantic/datafusion.git", branch = "pydantic-main" }

src/rewrite.rs

@@ -38,7 +38,7 @@ fn optimise_json_get_cast(cast: &Cast) -> Option<Transformed<Expr>> {
     if scalar_func.func.name() != "json_get" {
         return None;
     }
-    let func = match &cast.data_type {
+    let func = match cast.field.data_type() {
         DataType::Boolean => crate::json_get_bool::json_get_bool_udf(),
         DataType::Float64 | DataType::Float32 | DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => {
             crate::json_get_float::json_get_float_udf()

tests/main.rs

@@ -163,14 +163,27 @@ async fn test_json_get_array_with_path() {
 
 #[tokio::test]
 async fn test_json_get_equals() {
-    let e = run_query(r"select name, json_get(json_data, 'foo')='abc' from test")
+    // union comparison now works thanks to the union coercions upport in datafusion
+    // (previously failed with "Cannot infer common argument type for comparison operation Union")
+    // see https://github.com/apache/datafusion/issues/10180
+    let batches = run_query(r"select name, json_get(json_data, 'foo')='abc' from test")
         .await
-        .unwrap_err();
+        .unwrap();
 
-    // see https://github.com/apache/datafusion/issues/10180
-    assert!(e
-        .to_string()
-        .starts_with("Error during planning: Cannot infer common argument type for comparison operation Union"));
+    let expected = [
+        "+------------------+----------------------------------------------------+",
+        r#"| name             | json_get(test.json_data,Utf8("foo")) = Utf8("abc") |"#,
+        "+------------------+----------------------------------------------------+",
+        "| object_foo       | true                                               |",
+        "| object_foo_array |                                                    |",
+        "| object_foo_obj   |                                                    |",
+        "| object_foo_null  |                                                    |",
+        "| object_bar       |                                                    |",
+        "| list_foo         |                                                    |",
+        "| invalid_json     |                                                    |",
+        "+------------------+----------------------------------------------------+",
+    ];
+    assert_batches_eq!(expected, &batches);
 }
 
 #[tokio::test]