Don't return dictionary encoded booleans (#42)

samuelcolvin · adriangb · web-flow · commit b1ca610c0c83 · 2024-09-18T06:06:56.000+01:00
Co-authored-by: Adrian Garcia Badaracco &lt;1755071+adriangb@users.noreply.github.com&gt;
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -20,7 +20,7 @@ repos:
     pass_filenames: false
   - id: clippy
     name: Clippy
-    entry: cargo clippy
+    entry: cargo clippy -- -D warnings
     types: [rust]
     language: system
     pass_filenames: false
diff --git a/src/common.rs b/src/common.rs
@@ -5,7 +5,8 @@ use datafusion::arrow::array::{
     Array, ArrayRef, AsArray, DictionaryArray, Int64Array, LargeStringArray, PrimitiveArray, StringArray,
     StringViewArray, UInt64Array, UnionArray,
 };
-use datafusion::arrow::datatypes::{ArrowNativeType, ArrowPrimitiveType, DataType};
+use datafusion::arrow::compute::take;
+use datafusion::arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, ArrowPrimitiveType, DataType};
 use datafusion::arrow::downcast_dictionary_array;
 use datafusion::common::{exec_err, plan_err, Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
@@ -21,7 +22,7 @@ use crate::common_union::{is_json_union, json_from_union_scalar, nested_json_arr
 /// * `fn_name` - The name of the function
 /// * `value_type` - The general return type of the function, might be wrapped in a dictionary depending
 ///   on the first argument
-pub fn scalar_udf_return_type(args: &[DataType], fn_name: &str, value_type: DataType) -> DataFusionResult<DataType> {
+pub fn return_type_check(args: &[DataType], fn_name: &str, value_type: DataType) -> DataFusionResult<DataType> {
     let Some(first) = args.first() else {
         return plan_err!("The '{fn_name}' function requires one or more arguments.");
     };
@@ -105,6 +106,7 @@ pub fn invoke<C: FromIterator<Option<I>> + 'static, I>(
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
     to_array: impl Fn(C) -> DataFusionResult<ArrayRef>,
     to_scalar: impl Fn(Option<I>) -> ScalarValue,
+    return_dict: bool,
 ) -> DataFusionResult<ColumnarValue> {
     let Some(first_arg) = args.first() else {
         // I think this can't happen, but I assumed the same about args[1] and I was wrong, so better to be safe
@@ -118,13 +120,17 @@ pub fn invoke<C: FromIterator<Option<I>> + 'static, I>(
                         // TODO perhaps we could support this by zipping the arrays, but it's not trivial, #23
                         exec_err!("More than 1 path element is not supported when querying JSON using an array.")
                     } else {
-                        invoke_array(json_array, a, to_array, jiter_find)
+                        invoke_array(json_array, a, to_array, jiter_find, return_dict)
                     }
                 }
-                Some(ColumnarValue::Scalar(_)) => {
-                    scalar_apply(json_array, &JsonPath::extract_path(args), to_array, jiter_find)
-                }
-                None => scalar_apply(json_array, &[], to_array, jiter_find),
+                Some(ColumnarValue::Scalar(_)) => scalar_apply(
+                    json_array,
+                    &JsonPath::extract_path(args),
+                    to_array,
+                    jiter_find,
+                    return_dict,
+                ),
+                None => scalar_apply(json_array, &[], to_array, jiter_find, return_dict),
             };
             array.map(ColumnarValue::from)
         }
@@ -137,24 +143,26 @@ fn invoke_array<C: FromIterator<Option<I>> + 'static, I>(
     needle_array: &ArrayRef,
     to_array: impl Fn(C) -> DataFusionResult<ArrayRef>,
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
+    return_dict: bool,
 ) -> DataFusionResult<ArrayRef> {
     if let Some(d) = needle_array.as_any_dictionary_opt() {
-        invoke_array(json_array, d.values(), to_array, jiter_find)
+        // this is the (very rare) case where the needle is a dictionary, it shouldn't affect what we return
+        invoke_array(json_array, d.values(), to_array, jiter_find, return_dict)
     } else if let Some(str_path_array) = needle_array.as_any().downcast_ref::<StringArray>() {
         let paths = str_path_array.iter().map(|opt_key| opt_key.map(JsonPath::Key));
-        zip_apply(json_array, paths, to_array, jiter_find, true)
+        zip_apply(json_array, paths, to_array, jiter_find, true, return_dict)
     } else if let Some(str_path_array) = needle_array.as_any().downcast_ref::<LargeStringArray>() {
         let paths = str_path_array.iter().map(|opt_key| opt_key.map(JsonPath::Key));
-        zip_apply(json_array, paths, to_array, jiter_find, true)
+        zip_apply(json_array, paths, to_array, jiter_find, true, return_dict)
     } else if let Some(str_path_array) = needle_array.as_any().downcast_ref::<StringViewArray>() {
         let paths = str_path_array.iter().map(|opt_key| opt_key.map(JsonPath::Key));
-        zip_apply(json_array, paths, to_array, jiter_find, true)
+        zip_apply(json_array, paths, to_array, jiter_find, true, return_dict)
     } else if let Some(int_path_array) = needle_array.as_any().downcast_ref::<Int64Array>() {
         let paths = int_path_array.iter().map(|opt_index| opt_index.map(Into::into));
-        zip_apply(json_array, paths, to_array, jiter_find, false)
+        zip_apply(json_array, paths, to_array, jiter_find, false, return_dict)
     } else if let Some(int_path_array) = needle_array.as_any().downcast_ref::<UInt64Array>() {
         let paths = int_path_array.iter().map(|opt_index| opt_index.map(Into::into));
-        zip_apply(json_array, paths, to_array, jiter_find, false)
+        zip_apply(json_array, paths, to_array, jiter_find, false, return_dict)
     } else {
         exec_err!("unexpected second argument type, expected string or int array")
     }
@@ -166,22 +174,15 @@ fn zip_apply<'a, P: Iterator<Item = Option<JsonPath<'a>>>, C: FromIterator<Optio
     to_array: impl Fn(C) -> DataFusionResult<ArrayRef>,
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
     object_lookup: bool,
+    return_dict: bool,
 ) -> DataFusionResult<ArrayRef> {
     // arrow_schema "use" is workaround for https://github.com/apache/arrow-rs/issues/6400#issue-2528388332
     use datafusion::arrow::datatypes as arrow_schema;
 
     let c = downcast_dictionary_array!(
         json_array => {
-            let values = zip_apply(json_array.values(), path_array, to_array, jiter_find, object_lookup)?;
-            if !is_json_union(values.data_type()) {
-                return Ok(Arc::new(json_array.with_values(values)));
-            }
-            // JSON union: post-process the array to set keys to null where the union member is null
-            let type_ids = values.as_any().downcast_ref::<UnionArray>().unwrap().type_ids();
-            return Ok(Arc::new(DictionaryArray::new(
-                mask_dictionary_keys(json_array.keys(), type_ids),
-                values,
-            )));
+            let values = zip_apply(json_array.values(), path_array, to_array, jiter_find, object_lookup, false)?;
+            return post_process_dict(json_array, values, return_dict);
         }
         DataType::Utf8 => zip_apply_iter(json_array.as_string::<i32>().iter(), path_array, jiter_find),
         DataType::LargeUtf8 => zip_apply_iter(json_array.as_string::<i64>().iter(), path_array, jiter_find),
@@ -242,22 +243,15 @@ fn scalar_apply<C: FromIterator<Option<I>>, I>(
     path: &[JsonPath],
     to_array: impl Fn(C) -> DataFusionResult<ArrayRef>,
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
+    return_dict: bool,
 ) -> DataFusionResult<ArrayRef> {
     // arrow_schema "use" is workaround for https://github.com/apache/arrow-rs/issues/6400#issue-2528388332
     use datafusion::arrow::datatypes as arrow_schema;
 
     let c = downcast_dictionary_array!(
         json_array => {
-            let values = scalar_apply(json_array.values(), path, to_array, jiter_find)?;
-            if !is_json_union(values.data_type()) {
-                return Ok(Arc::new(json_array.with_values(values)));
-            }
-            // JSON union: post-process the array to set keys to null where the union member is null
-            let type_ids = values.as_any().downcast_ref::<UnionArray>().unwrap().type_ids();
-            return Ok(Arc::new(DictionaryArray::new(
-                mask_dictionary_keys(json_array.keys(), type_ids),
-                values,
-            )));
+            let values = scalar_apply(json_array.values(), path, to_array, jiter_find, false)?;
+            return post_process_dict(json_array, values, return_dict);
         }
         DataType::Utf8 => scalar_apply_iter(json_array.as_string::<i32>().iter(), path, jiter_find),
         DataType::LargeUtf8 => scalar_apply_iter(json_array.as_string::<i64>().iter(), path, jiter_find),
@@ -268,10 +262,32 @@ fn scalar_apply<C: FromIterator<Option<I>>, I>(
             return exec_err!("unexpected json array type {:?}", other);
         }
     );
-
     to_array(c)
 }
 
+/// Take a dictionary array of JSON data and an array of result values and combine them.
+fn post_process_dict<T: ArrowDictionaryKeyType>(
+    dict_array: &DictionaryArray<T>,
+    result_values: ArrayRef,
+    return_dict: bool,
+) -> DataFusionResult<ArrayRef> {
+    if return_dict {
+        if is_json_union(result_values.data_type()) {
+            // JSON union: post-process the array to set keys to null where the union member is null
+            let type_ids = result_values.as_any().downcast_ref::<UnionArray>().unwrap().type_ids();
+            Ok(Arc::new(DictionaryArray::new(
+                mask_dictionary_keys(dict_array.keys(), type_ids),
+                result_values,
+            )))
+        } else {
+            Ok(Arc::new(dict_array.with_values(result_values)))
+        }
+    } else {
+        // this is what cast would do under the hood to unpack a dictionary into an array of its values
+        Ok(take(&result_values, dict_array.keys(), None)?)
+    }
+}
+
 fn is_object_lookup(path: &[JsonPath]) -> bool {
     if let Some(first) = path.first() {
         matches!(first, JsonPath::Key(_))
diff --git a/src/json_as_text.rs b/src/json_as_text.rs
@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::Peek;
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonAsText {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Utf8)
+        return_type_check(arg_types, self.name(), DataType::Utf8)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonAsText {
             jiter_json_as_text,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Utf8,
+            true,
         )
     }
 
diff --git a/src/json_contains.rs b/src/json_contains.rs
@@ -6,7 +6,7 @@ use datafusion::common::arrow::array::{ArrayRef, BooleanArray};
 use datafusion::common::{plan_err, Result, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 
-use crate::common::{invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -48,7 +48,7 @@ impl ScalarUDFImpl for JsonContains {
         if arg_types.len() < 2 {
             plan_err!("The 'json_contains' function requires two or more arguments.")
         } else {
-            scalar_udf_return_type(arg_types, self.name(), DataType::Boolean)
+            return_type_check(arg_types, self.name(), DataType::Boolean).map(|_| DataType::Boolean)
         }
     }
 
@@ -58,6 +58,7 @@ impl ScalarUDFImpl for JsonContains {
             jiter_json_contains,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Boolean,
+            false,
         )
     }
 
diff --git a/src/json_get.rs b/src/json_get.rs
@@ -8,7 +8,7 @@ use datafusion::common::Result as DataFusionResult;
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::{Jiter, NumberAny, NumberInt, Peek};
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 use crate::common_union::{JsonUnion, JsonUnionField};
 
@@ -50,15 +50,15 @@ impl ScalarUDFImpl for JsonGet {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), JsonUnion::data_type())
+        return_type_check(arg_types, self.name(), JsonUnion::data_type())
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
         let to_array = |c: JsonUnion| {
             let array: UnionArray = c.try_into()?;
             Ok(Arc::new(array) as ArrayRef)
         };
-        invoke::<JsonUnion, JsonUnionField>(args, jiter_json_get_union, to_array, JsonUnionField::scalar_value)
+        invoke::<JsonUnion, JsonUnionField>(args, jiter_json_get_union, to_array, JsonUnionField::scalar_value, true)
     }
 
     fn aliases(&self) -> &[String] {
diff --git a/src/json_get_bool.rs b/src/json_get_bool.rs
@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::Peek;
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetBool {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Boolean)
+        return_type_check(arg_types, self.name(), DataType::Boolean).map(|_| DataType::Boolean)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetBool {
             jiter_json_get_bool,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Boolean,
+            false,
         )
     }
 
diff --git a/src/json_get_float.rs b/src/json_get_float.rs
@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::{NumberAny, Peek};
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetFloat {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Float64)
+        return_type_check(arg_types, self.name(), DataType::Float64)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetFloat {
             jiter_json_get_float,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Float64,
+            true,
         )
     }
 
diff --git a/src/json_get_int.rs b/src/json_get_int.rs
@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::{NumberInt, Peek};
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetInt {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Int64)
+        return_type_check(arg_types, self.name(), DataType::Int64)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetInt {
             jiter_json_get_int,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Int64,
+            true,
         )
     }
 
diff --git a/src/json_get_json.rs b/src/json_get_json.rs
@@ -6,7 +6,7 @@ use datafusion::arrow::datatypes::DataType;
 use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -45,7 +45,7 @@ impl ScalarUDFImpl for JsonGetJson {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Utf8)
+        return_type_check(arg_types, self.name(), DataType::Utf8)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -54,6 +54,7 @@ impl ScalarUDFImpl for JsonGetJson {
             jiter_json_get_json,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Utf8,
+            true,
         )
     }
 
diff --git a/src/json_get_str.rs b/src/json_get_str.rs
@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use jiter::Peek;
 
-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};
+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetStr {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
-        scalar_udf_return_type(arg_types, self.name(), DataType::Utf8)
+        return_type_check(arg_types, self.name(), DataType::Utf8)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetStr {
             jiter_json_get_str,
             |c| Ok(Arc::new(c) as ArrayRef),
             ScalarValue::Utf8,
+            true,
         )
     }
 
diff --git a/src/json_length.rs b/src/json_length.rs
diff --git a/tests/main.rs b/tests/main.rs
diff --git a/tests/utils/mod.rs b/tests/utils/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};`
`7`	`7`	`use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};`
`8`	`8`	`use jiter::Peek;`
`9`	`9`
`10`		`-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};`
	`10`	`+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};`
`11`	`11`	`use crate::common_macros::make_udf_function;`
`12`	`12`
`13`	`13`	`make_udf_function!(`
`@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonAsText {`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {`
`49`		`- scalar_udf_return_type(arg_types, self.name(), DataType::Utf8)`
	`49`	`+ return_type_check(arg_types, self.name(), DataType::Utf8)`
`50`	`50`	`}`
`51`	`51`
`52`	`52`	`fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {`
`@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonAsText {`
`55`	`55`	`jiter_json_as_text,`
`56`	`56`	`\|c\| Ok(Arc::new(c) as ArrayRef),`
`57`	`57`	`ScalarValue::Utf8,`
	`58`	`+ true,`
`58`	`59`	`)`
`59`	`60`	`}`
`60`	`61`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ use datafusion::common::arrow::array::{ArrayRef, BooleanArray};`
`6`	`6`	`use datafusion::common::{plan_err, Result, ScalarValue};`
`7`	`7`	`use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};`
`8`	`8`
`9`		`-use crate::common::{invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};`
	`9`	`+use crate::common::{invoke, jiter_json_find, return_type_check, GetError, JsonPath};`
`10`	`10`	`use crate::common_macros::make_udf_function;`
`11`	`11`
`12`	`12`	`make_udf_function!(`
`@@ -48,7 +48,7 @@ impl ScalarUDFImpl for JsonContains {`
`48`	`48`	`if arg_types.len() < 2 {`
`49`	`49`	`plan_err!("The 'json_contains' function requires two or more arguments.")`
`50`	`50`	`} else {`
`51`		`- scalar_udf_return_type(arg_types, self.name(), DataType::Boolean)`
	`51`	`+ return_type_check(arg_types, self.name(), DataType::Boolean).map(\|_\| DataType::Boolean)`
`52`	`52`	`}`
`53`	`53`	`}`
`54`	`54`
`@@ -58,6 +58,7 @@ impl ScalarUDFImpl for JsonContains {`
`58`	`58`	`jiter_json_contains,`
`59`	`59`	`\|c\| Ok(Arc::new(c) as ArrayRef),`
`60`	`60`	`ScalarValue::Boolean,`
	`61`	`+ false,`
`61`	`62`	`)`
`62`	`63`	`}`
`63`	`64`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};`
`7`	`7`	`use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};`
`8`	`8`	`use jiter::{NumberAny, Peek};`
`9`	`9`
`10`		`-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};`
	`10`	`+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};`
`11`	`11`	`use crate::common_macros::make_udf_function;`
`12`	`12`
`13`	`13`	`make_udf_function!(`
`@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetFloat {`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {`
`49`		`- scalar_udf_return_type(arg_types, self.name(), DataType::Float64)`
	`49`	`+ return_type_check(arg_types, self.name(), DataType::Float64)`
`50`	`50`	`}`
`51`	`51`
`52`	`52`	`fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {`
`@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetFloat {`
`55`	`55`	`jiter_json_get_float,`
`56`	`56`	`\|c\| Ok(Arc::new(c) as ArrayRef),`
`57`	`57`	`ScalarValue::Float64,`
	`58`	`+ true,`
`58`	`59`	`)`
`59`	`60`	`}`
`60`	`61`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ use datafusion::common::{Result as DataFusionResult, ScalarValue};`
`7`	`7`	`use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};`
`8`	`8`	`use jiter::{NumberInt, Peek};`
`9`	`9`
`10`		`-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};`
	`10`	`+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};`
`11`	`11`	`use crate::common_macros::make_udf_function;`
`12`	`12`
`13`	`13`	`make_udf_function!(`
`@@ -46,7 +46,7 @@ impl ScalarUDFImpl for JsonGetInt {`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {`
`49`		`- scalar_udf_return_type(arg_types, self.name(), DataType::Int64)`
	`49`	`+ return_type_check(arg_types, self.name(), DataType::Int64)`
`50`	`50`	`}`
`51`	`51`
`52`	`52`	`fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {`
`@@ -55,6 +55,7 @@ impl ScalarUDFImpl for JsonGetInt {`
`55`	`55`	`jiter_json_get_int,`
`56`	`56`	`\|c\| Ok(Arc::new(c) as ArrayRef),`
`57`	`57`	`ScalarValue::Int64,`
	`58`	`+ true,`
`58`	`59`	`)`
`59`	`60`	`}`
`60`	`61`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ use datafusion::arrow::datatypes::DataType;`
`6`	`6`	`use datafusion::common::{Result as DataFusionResult, ScalarValue};`
`7`	`7`	`use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};`
`8`	`8`
`9`		`-use crate::common::{get_err, invoke, jiter_json_find, scalar_udf_return_type, GetError, JsonPath};`
	`9`	`+use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};`
`10`	`10`	`use crate::common_macros::make_udf_function;`
`11`	`11`
`12`	`12`	`make_udf_function!(`
`@@ -45,7 +45,7 @@ impl ScalarUDFImpl for JsonGetJson {`
`45`	`45`	`}`
`46`	`46`
`47`	`47`	`fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {`
`48`		`- scalar_udf_return_type(arg_types, self.name(), DataType::Utf8)`
	`48`	`+ return_type_check(arg_types, self.name(), DataType::Utf8)`
`49`	`49`	`}`
`50`	`50`
`51`	`51`	`fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {`
`@@ -54,6 +54,7 @@ impl ScalarUDFImpl for JsonGetJson {`
`54`	`54`	`jiter_json_get_json,`
`55`	`55`	`\|c\| Ok(Arc::new(c) as ArrayRef),`
`56`	`56`	`ScalarValue::Utf8,`
	`57`	`+ true,`
`57`	`58`	`)`
`58`	`59`	`}`
`59`	`60`