support function unnesting for json_as_text inner calls (#57)

Author: davidhewitt

src/rewrite.rs

@@ -52,14 +52,23 @@ fn optimise_json_get_cast(cast: &Cast) -> Option<Transformed<Expr>> {
 fn unnest_json_calls(func: &ScalarFunction) -> Option<Transformed<Expr>> {
     if !matches!(
         func.func.name(),
-        "json_get" | "json_get_bool" | "json_get_float" | "json_get_int" | "json_get_json" | "json_get_str"
+        "json_get"
+            | "json_get_bool"
+            | "json_get_float"
+            | "json_get_int"
+            | "json_get_json"
+            | "json_get_str"
+            | "json_as_text"
     ) {
         return None;
     }
     let mut outer_args_iter = func.args.iter();
     let first_arg = outer_args_iter.next()?;
     let inner_func = extract_scalar_function(first_arg)?;
-    if inner_func.func.name() != "json_get" {
+
+    // both json_get and json_as_text would produce new JSON to be processed by the outer
+    // function so can be inlined
+    if !matches!(inner_func.func.name(), "json_get" | "json_as_text") {
         return None;
     }
 

tests/main.rs

@@ -880,6 +880,38 @@ async fn test_plan_arrow_double_nested() {
     assert_eq!(lines, expected);
 }
 
+#[tokio::test]
+async fn test_double_arrow_double_nested() {
+    let batches = run_query("select name, json_data->>'foo'->>0 from test").await.unwrap();
+
+    let expected = [
+        "+------------------+---------------------------------------------+",
+        "| name             | test.json_data ->> Utf8(\"foo\") ->> Int64(0) |",
+        "+------------------+---------------------------------------------+",
+        "| object_foo       |                                             |",
+        "| object_foo_array | 1                                           |",
+        "| object_foo_obj   |                                             |",
+        "| object_foo_null  |                                             |",
+        "| object_bar       |                                             |",
+        "| list_foo         |                                             |",
+        "| invalid_json     |                                             |",
+        "+------------------+---------------------------------------------+",
+    ];
+    assert_batches_eq!(expected, &batches);
+}
+
+#[tokio::test]
+async fn test_plan_double_arrow_double_nested() {
+    let lines = logical_plan(r"explain select json_data->>'foo'->>0 from test").await;
+
+    let expected = [
+        "Projection: json_as_text(test.json_data, Utf8(\"foo\"), Int64(0)) AS test.json_data ->> Utf8(\"foo\") ->> Int64(0)",
+        "  TableScan: test projection=[json_data]",
+    ];
+
+    assert_eq!(lines, expected);
+}
+
 #[tokio::test]
 async fn test_arrow_double_nested_cast() {
     let batches = run_query("select name, (json_data->'foo'->0)::int from test")
@@ -914,6 +946,41 @@ async fn test_plan_arrow_double_nested_cast() {
     assert_eq!(lines, expected);
 }
 
+#[tokio::test]
+async fn test_double_arrow_double_nested_cast() {
+    let batches = run_query("select name, (json_data->>'foo'->>0)::int from test")
+        .await
+        .unwrap();
+
+    let expected = [
+        "+------------------+---------------------------------------------+",
+        "| name             | test.json_data ->> Utf8(\"foo\") ->> Int64(0) |",
+        "+------------------+---------------------------------------------+",
+        "| object_foo       |                                             |",
+        "| object_foo_array | 1                                           |",
+        "| object_foo_obj   |                                             |",
+        "| object_foo_null  |                                             |",
+        "| object_bar       |                                             |",
+        "| list_foo         |                                             |",
+        "| invalid_json     |                                             |",
+        "+------------------+---------------------------------------------+",
+    ];
+    assert_batches_eq!(expected, &batches);
+}
+
+#[tokio::test]
+async fn test_plan_double_arrow_double_nested_cast() {
+    let lines = logical_plan(r"explain select (json_data->>'foo'->>0)::int from test").await;
+
+    // NB: json_as_text(..)::int is NOT the same as `json_get_int(..)`, hence the cast is not rewritten
+    let expected = [
+        "Projection: CAST(json_as_text(test.json_data, Utf8(\"foo\"), Int64(0)) AS test.json_data ->> Utf8(\"foo\") ->> Int64(0) AS Int32)",
+        "  TableScan: test projection=[json_data]",
+    ];
+
+    assert_eq!(lines, expected);
+}
+
 #[tokio::test]
 async fn test_arrow_nested_columns() {
     let expected = [