unnest json_get calls (#25)

Author: samuelcolvin

Cargo.toml

@@ -8,7 +8,7 @@ license = "Apache-2.0"
 keywords = ["datafusion", "JSON", "SQL"]
 categories = ["database-implementations", "parsing"]
 repository = "https://github.com/datafusion-contrib/datafusion-functions-json/"
-rust-version = "1.73.0"
+rust-version = "1.76.0"
 
 [dependencies]
 arrow = "52"
@@ -24,7 +24,7 @@ datafusion-execution = "39"
 codspeed-criterion-compat = "2.3"
 criterion = "0.5.1"
 datafusion = "39"
-clap = "~4.4" # for testing on MSRV 1.73
+clap = "4"
 tokio = { version = "1.37", features = ["full"] }
 
 [lints.clippy]

src/common.rs

@@ -75,6 +75,12 @@ pub fn invoke<C: FromIterator<Option<I>> + 'static, I>(
         ColumnarValue::Array(json_array) => {
             let result_collect = match args.get(1) {
                 Some(ColumnarValue::Array(a)) => {
+                    if args.len() > 2 {
+                        // TODO perhaps we could support this by zipping the arrays, but it's not trivial, #23
+                        return exec_err!(
+                            "More than 1 path element is not supported when querying JSON using an array."
+                        );
+                    }
                     if let Some(str_path_array) = a.as_any().downcast_ref::<StringArray>() {
                         let paths = str_path_array.iter().map(|opt_key| opt_key.map(JsonPath::Key));
                         zip_apply(json_array, paths, jiter_find, true)
@@ -114,28 +120,28 @@ pub fn invoke<C: FromIterator<Option<I>> + 'static, I>(
 
 fn zip_apply<'a, P: Iterator<Item = Option<JsonPath<'a>>>, C: FromIterator<Option<I>> + 'static, I>(
     json_array: &ArrayRef,
-    paths: P,
+    path_array: P,
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
     object_lookup: bool,
 ) -> DataFusionResult<C> {
     if let Some(string_array) = json_array.as_any().downcast_ref::<StringArray>() {
-        Ok(zip_apply_iter(string_array.iter(), paths, jiter_find))
+        Ok(zip_apply_iter(string_array.iter(), path_array, jiter_find))
     } else if let Some(large_string_array) = json_array.as_any().downcast_ref::<LargeStringArray>() {
-        Ok(zip_apply_iter(large_string_array.iter(), paths, jiter_find))
+        Ok(zip_apply_iter(large_string_array.iter(), path_array, jiter_find))
     } else if let Some(string_array) = nested_json_array(json_array, object_lookup) {
-        Ok(zip_apply_iter(string_array.iter(), paths, jiter_find))
+        Ok(zip_apply_iter(string_array.iter(), path_array, jiter_find))
     } else {
         exec_err!("unexpected json array type {:?}", json_array.data_type())
     }
 }
 
 fn zip_apply_iter<'a, 'j, P: Iterator<Item = Option<JsonPath<'a>>>, C: FromIterator<Option<I>> + 'static, I>(
     json_iter: impl Iterator<Item = Option<&'j str>>,
-    paths: P,
+    path_array: P,
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
 ) -> C {
     json_iter
-        .zip(paths)
+        .zip(path_array)
         .map(|(opt_json, opt_path)| {
             if let Some(path) = opt_path {
                 jiter_find(opt_json, &[path]).ok()

src/rewrite.rs

@@ -23,11 +23,45 @@ impl FunctionRewrite for JsonFunctionRewriter {
                     }
                 }
             }
+        } else if let Expr::ScalarFunction(func) = &expr {
+            if let Some(new_func) = unnest_json_calls(func) {
+                return Ok(Transformed::yes(Expr::ScalarFunction(new_func)));
+            }
         }
         Ok(Transformed::no(expr))
     }
 }
 
+// Replace nested JSON functions e.g. `json_get(json_get(col, 'foo'), 'bar')` with `json_get(col, 'foo', 'bar')`
+fn unnest_json_calls(func: &ScalarFunction) -> Option<ScalarFunction> {
+    if !matches!(
+        func.func.name(),
+        "json_get" | "json_get_bool" | "json_get_float" | "json_get_int" | "json_get_json" | "json_get_str"
+    ) {
+        return None;
+    }
+    let mut outer_args_iter = func.args.iter();
+    let first_arg = outer_args_iter.next()?;
+    let Expr::ScalarFunction(inner_func) = first_arg else {
+        return None;
+    };
+    if inner_func.func.name() != "json_get" {
+        return None;
+    }
+
+    let mut args = inner_func.args.clone();
+    args.extend(outer_args_iter.cloned());
+    // See #23, unnest only when all lookup arguments are literals
+    if args.iter().skip(1).all(|arg| matches!(arg, Expr::Literal(_))) {
+        Some(ScalarFunction {
+            func: func.func.clone(),
+            args,
+        })
+    } else {
+        None
+    }
+}
+
 fn switch_json_get(cast_data_type: &DataType, args: &[Expr]) -> Option<Transformed<Expr>> {
     let func = match cast_data_type {
         DataType::Boolean => crate::json_get_bool::json_get_bool_udf(),

tests/main.rs

@@ -5,7 +5,7 @@ use datafusion_common::ScalarValue;
 mod utils;
 use datafusion_expr::ColumnarValue;
 use datafusion_functions_json::udfs::json_get_str_udf;
-use utils::{display_val, run_query, run_query_large, run_query_params};
+use utils::{display_val, logical_plan, run_query, run_query_large, run_query_params};
 
 #[tokio::test]
 async fn test_json_contains() {
@@ -523,7 +523,7 @@ async fn test_json_get_union_scalar() {
 }
 
 #[tokio::test]
-async fn test_json_get_union_array() {
+async fn test_json_get_nested_collapsed() {
     let expected = [
         "+------------------+---------+",
         "| name             | v       |",
@@ -545,7 +545,130 @@ async fn test_json_get_union_array() {
 }
 
 #[tokio::test]
-async fn test_json_get_union_array_skip() {
+async fn test_json_get_cte() {
+    // avoid auto-un-nesting with a CTE
+    let sql = r#"
+        with t as (select name, json_get(json_data, 'foo') j from test)
+        select name, json_get(j, 0) v from t
+    "#;
+    let expected = [
+        "+------------------+---------+",
+        "| name             | v       |",
+        "+------------------+---------+",
+        "| object_foo       | {null=} |",
+        "| object_foo_array | {int=1} |",
+        "| object_foo_obj   | {null=} |",
+        "| object_foo_null  | {null=} |",
+        "| object_bar       | {null=} |",
+        "| list_foo         | {null=} |",
+        "| invalid_json     | {null=} |",
+        "+------------------+---------+",
+    ];
+
+    let batches = run_query(sql).await.unwrap();
+    assert_batches_eq!(expected, &batches);
+}
+
+#[tokio::test]
+async fn test_json_get_cte_plan() {
+    // avoid auto-unnesting with a CTE
+    let sql = r#"
+        explain
+        with t as (select name, json_get(json_data, 'foo') j from test)
+        select name, json_get(j, 0) v from t
+    "#;
+    let expected = [
+        "Projection: t.name, json_get(t.j, Int64(0)) AS v",
+        "  SubqueryAlias: t",
+        "    Projection: test.name, json_get(test.json_data, Utf8(\"foo\")) AS j",
+        "      TableScan: test projection=[name, json_data]",
+    ];
+
+    let plan_lines = logical_plan(sql).await;
+    assert_eq!(plan_lines, expected);
+}
+
+#[tokio::test]
+async fn test_json_get_unnest() {
+    let sql = "select name, json_get(json_get(json_data, 'foo'), 0) v from test";
+
+    let expected = [
+        "+------------------+---------+",
+        "| name             | v       |",
+        "+------------------+---------+",
+        "| object_foo       | {null=} |",
+        "| object_foo_array | {int=1} |",
+        "| object_foo_obj   | {null=} |",
+        "| object_foo_null  | {null=} |",
+        "| object_bar       | {null=} |",
+        "| list_foo         | {null=} |",
+        "| invalid_json     | {null=} |",
+        "+------------------+---------+",
+    ];
+
+    let batches = run_query(sql).await.unwrap();
+    assert_batches_eq!(expected, &batches);
+}
+
+#[tokio::test]
+async fn test_json_get_unnest_plan() {
+    let sql = "explain select json_get(json_get(json_data, 'foo'), 0) v from test";
+    let expected = [
+        "Projection: json_get(test.json_data, Utf8(\"foo\"), Int64(0)) AS v",
+        "  TableScan: test projection=[json_data]",
+    ];
+
+    let plan_lines = logical_plan(sql).await;
+    assert_eq!(plan_lines, expected);
+}
+
+#[tokio::test]
+async fn test_json_get_int_unnest() {
+    let sql = "select name, json_get(json_get(json_data, 'foo'), 0)::int v from test";
+
+    let expected = [
+        "+------------------+---+",
+        "| name             | v |",
+        "+------------------+---+",
+        "| object_foo       |   |",
+        "| object_foo_array | 1 |",
+        "| object_foo_obj   |   |",
+        "| object_foo_null  |   |",
+        "| object_bar       |   |",
+        "| list_foo         |   |",
+        "| invalid_json     |   |",
+        "+------------------+---+",
+    ];
+
+    let batches = run_query(sql).await.unwrap();
+    assert_batches_eq!(expected, &batches);
+}
+
+#[tokio::test]
+async fn test_json_get_int_unnest_plan() {
+    let sql = "explain select json_get(json_get(json_data, 'foo'), 0)::int v from test";
+    let expected = [
+        "Projection: json_get_int(test.json_data, Utf8(\"foo\"), Int64(0)) AS v",
+        "  TableScan: test projection=[json_data]",
+    ];
+
+    let plan_lines = logical_plan(sql).await;
+    assert_eq!(plan_lines, expected);
+}
+
+#[tokio::test]
+async fn test_multiple_lookup_arrays() {
+    let sql = "select json_get(json_data, str_key1, str_key2) v from more_nested";
+    let err = run_query(sql).await.unwrap_err();
+    assert_eq!(
+        err.to_string(),
+        "Execution error: More than 1 path element is not supported when querying JSON using an array."
+    );
+}
+
+#[tokio::test]
+async fn test_json_get_union_array_nested() {
+    let sql = "select json_get(json_get(json_data, str_key1), str_key2) v from more_nested";
     let expected = [
         "+-------------+",
         "| v           |",
@@ -556,13 +679,27 @@ async fn test_json_get_union_array_skip() {
         "+-------------+",
     ];
 
-    let sql = "select json_get(json_get(json_data, str_key1), str_key2) v from more_nested";
     let batches = run_query(sql).await.unwrap();
     assert_batches_eq!(expected, &batches);
 }
 
+#[tokio::test]
+async fn test_json_get_union_array_nested_plan() {
+    let sql = "explain select json_get(json_get(json_data, str_key1), str_key2) v from more_nested";
+    // json_get is not un-nested because lookup types are not literals
+    let expected = [
+        "Projection: json_get(json_get(more_nested.json_data, more_nested.str_key1), more_nested.str_key2) AS v",
+        "  TableScan: more_nested projection=[json_data, str_key1, str_key2]",
+    ];
+
+    let plan_lines = logical_plan(sql).await;
+    assert_eq!(plan_lines, expected);
+}
+
 #[tokio::test]
 async fn test_json_get_union_array_skip_double_nested() {
+    let sql =
+        "select json_data, json_get_int(json_get(json_get(json_data, str_key1), str_key2), int_key) v from more_nested";
     let expected = [
         "+--------------------------+---+",
         "| json_data                | v |",
@@ -573,8 +710,6 @@ async fn test_json_get_union_array_skip_double_nested() {
         "+--------------------------+---+",
     ];
 
-    let sql =
-        "select json_data, json_get_int(json_get(json_get(json_data, str_key1), str_key2), int_key) v from more_nested";
     let batches = run_query(sql).await.unwrap();
     assert_batches_eq!(expected, &batches);
 }

tests/utils/mod.rs

@@ -142,3 +142,10 @@ pub async fn display_val(batch: Vec<RecordBatch>) -> (DataType, String) {
     let repr = f.value(0).try_to_string().unwrap();
     (schema_col.data_type().clone(), repr)
 }
+
+pub async fn logical_plan(sql: &str) -> Vec<String> {
+    let batches = run_query(sql).await.unwrap();
+    let plan_col = batches[0].column(1).as_any().downcast_ref::<StringArray>().unwrap();
+    let logical_plan = plan_col.value(0);
+    logical_plan.split('\n').map(|s| s.to_string()).collect()
+}