add failing test_json_get_equals test

Cargo.toml

@@ -25,3 +25,4 @@ print_stdout = "warn"
 # certain lints which we don't want to enforce (for now)
 pedantic = { level = "warn", priority = -1 }
 missing_errors_doc = "allow"
+cast_possible_truncation = "allow"

src/json_get.rs

@@ -82,7 +82,6 @@ impl ScalarUDFImpl for JsonGet {
         let mut union = JsonUnion::new(json_data.len());
         for opt_json in json_data {
             if let Some(union_field) = jiter_json_get(opt_json, &path) {
-                dbg!(&union_field);
                 union.push(union_field);
             } else {
                 union.push_none();
@@ -126,11 +125,11 @@ fn _jiter_json_get(jiter: &mut Jiter, peek: Peek, path: &[JsonPath]) -> Result<J
     let next_peek = match peek {
         Peek::Array => match first {
             JsonPath::Index(index) => jiter_array_get(jiter, *index),
-            _ => Err(GetError),
+            JsonPath::Key(_) => Err(GetError),
         },
         Peek::Object => match first {
             JsonPath::Key(key) => jiter_object_get(jiter, key),
-            _ => Err(GetError),
+            JsonPath::Index(_) => Err(GetError),
         },
         _ => Err(GetError),
     }?;

src/lib.rs

@@ -7,6 +7,7 @@ use std::sync::Arc;
 mod json_get;
 mod json_obj_contains;
 mod macros;
+mod rewrite;
 mod union;
 
 pub mod functions {
@@ -24,6 +25,7 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
         }
         Ok(()) as Result<()>
     })?;
+    registry.register_function_rewrite(Arc::new(rewrite::JsonFunctionRewriter))?;
 
     Ok(())
 }

src/rewrite.rs

@@ -0,0 +1,19 @@
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::DFSchema;
+use datafusion_common::Result;
+use datafusion_expr::expr_rewriter::FunctionRewrite;
+use datafusion_expr::Expr;
+
+pub(crate) struct JsonFunctionRewriter;
+
+impl FunctionRewrite for JsonFunctionRewriter {
+    fn name(&self) -> &str {
+        "JsonFunctionRewriter"
+    }
+
+    fn rewrite(&self, expr: Expr, _schema: &DFSchema, _config: &ConfigOptions) -> Result<Transformed<Expr>> {
+        dbg!(&expr, _schema);
+        Ok(Transformed::no(expr))
+    }
+}

tests/test_json_get.rs

@@ -24,3 +24,25 @@ async fn test_json_get() {
         .unwrap();
     assert_batches_eq!(expected, &batches);
 }
+
+#[tokio::test]
+async fn test_json_get_equals() {
+    let expected = [
+        "+------------------+--------------------------------------+",
+        "| name             | json_get(test.json_data,Utf8(\"foo\")) |",
+        "+------------------+--------------------------------------+",
+        "| object_foo       | {int=123}                            |",
+        "| object_foo_array | {array=[1]}                          |",
+        "| object_foo_obj   | {object={}}                          |",
+        "| object_foo_null  | {null=true}                          |",
+        "| object_bar       | {null=}                              |",
+        "| list_foo         | {null=}                              |",
+        "| invalid_json     | {null=}                              |",
+        "+------------------+--------------------------------------+",
+    ];
+
+    let batches = run_query("select name, json_get(json_data, 'foo')=123 from test")
+        .await
+        .unwrap();
+    assert_batches_eq!(expected, &batches);
+}