fix optimizations

Author: adriangb

src/common.rs

@@ -612,6 +612,14 @@ pub(crate) enum Sortedness {
     Recursive,
 }
 
+impl Sortedness {
+    pub(crate) fn iter() -> impl Iterator<Item = Self> {
+        [Sortedness::Unspecified, Sortedness::TopLevel, Sortedness::Recursive]
+            .iter()
+            .copied()
+    }
+}
+
 impl Sortedness {
     pub(crate) fn function_name_suffix(self) -> &'static str {
         match self {

src/rewrite.rs

@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::sync::LazyLock;
 
 use datafusion::arrow::datatypes::DataType;
 use datafusion::common::config::ConfigOptions;
@@ -11,8 +12,12 @@ use datafusion::logical_expr::expr_rewriter::FunctionRewrite;
 use datafusion::logical_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
 use datafusion::logical_expr::sqlparser::ast::BinaryOperator;
 use datafusion::logical_expr::ScalarUDF;
+use datafusion::logical_expr::ScalarUDFImpl;
 use datafusion::scalar::ScalarValue;
 
+use crate::common::Sortedness;
+use crate::json_get::JsonGet;
+
 #[derive(Debug)]
 pub(crate) struct JsonFunctionRewriter;
 
@@ -31,11 +36,14 @@ impl FunctionRewrite for JsonFunctionRewriter {
     }
 }
 
+static JSON_GET_FUNC_NAMES: LazyLock<Vec<String>> =
+    LazyLock::new(|| Sortedness::iter().map(|s| JsonGet::new(s).name().to_string()).collect());
+
 /// This replaces `get_json(foo, bar)::int` with `json_get_int(foo, bar)` so the JSON function can take care of
 /// extracting the right value type from JSON without the need to materialize the JSON union.
 fn optimise_json_get_cast(cast: &Cast) -> Option<Transformed<Expr>> {
     let scalar_func = extract_scalar_function(&cast.expr)?;
-    if scalar_func.func.name() != "json_get" {
+    if !JSON_GET_FUNC_NAMES.contains(&scalar_func.func.name().to_owned()) {
         return None;
     }
     let func = match &cast.data_type {
@@ -53,18 +61,24 @@ fn optimise_json_get_cast(cast: &Cast) -> Option<Transformed<Expr>> {
     })))
 }
 
+static JSON_FUNCTION_NAMES: LazyLock<Vec<String>> = LazyLock::new(|| {
+    Sortedness::iter()
+        .flat_map(|s| {
+            [
+                crate::json_get::JsonGet::new(s).name().to_string(),
+                crate::json_get_bool::JsonGetBool::new(s).name().to_string(),
+                crate::json_get_float::JsonGetFloat::new(s).name().to_string(),
+                crate::json_get_int::JsonGetInt::new(s).name().to_string(),
+                crate::json_get_str::JsonGetStr::new(s).name().to_string(),
+                crate::json_as_text::JsonAsText::new(s).name().to_string(),
+            ]
+        })
+        .collect()
+});
+
 // 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<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_as_text"
-    ) {
+    if !JSON_FUNCTION_NAMES.contains(&func.func.name().to_owned()) {
         return None;
     }
     let mut outer_args_iter = func.args.iter();

tests/main.rs

@@ -346,6 +346,10 @@ async fn test_json_get_cast_int() {
     let batches = run_query(sql).await.unwrap();
     assert_eq!(display_val(batches).await, (DataType::Int64, "42".to_string()));
 
+    let sql = r#"select json_get_top_level_sorted('{"foo": 42}', 'foo')::int"#;
+    let batches = run_query(sql).await.unwrap();
+    assert_eq!(display_val(batches).await, (DataType::Int64, "42".to_string()));
+
     // floats not allowed
     let sql = r#"select json_get('{"foo": 4.2}', 'foo')::int"#;
     let batches = run_query(sql).await.unwrap();
@@ -400,13 +404,21 @@ async fn test_json_get_cast_float() {
     let sql = r#"select json_get('{"foo": 4.2e2}', 'foo')::float"#;
     let batches = run_query(sql).await.unwrap();
     assert_eq!(display_val(batches).await, (DataType::Float64, "420.0".to_string()));
+
+    let sql = r#"select json_get_top_level_sorted('{"foo": 4.2e2}', 'foo')::float"#;
+    let batches = run_query(sql).await.unwrap();
+    assert_eq!(display_val(batches).await, (DataType::Float64, "420.0".to_string()));
 }
 
 #[tokio::test]
 async fn test_json_get_cast_numeric() {
     let sql = r#"select json_get('{"foo": 4.2e2}', 'foo')::numeric"#;
     let batches = run_query(sql).await.unwrap();
     assert_eq!(display_val(batches).await, (DataType::Float64, "420.0".to_string()));
+
+    let sql = r#"select json_get_top_level_sorted('{"foo": 4.2e2}', 'foo')::numeric"#;
+    let batches = run_query(sql).await.unwrap();
+    assert_eq!(display_val(batches).await, (DataType::Float64, "420.0".to_string()));
 }
 
 #[tokio::test]