add variants of functions optimized for sorted keys

Author: adriangb

benches/main.rs

@@ -2,7 +2,7 @@ use codspeed_criterion_compat::{criterion_group, criterion_main, Bencher, Criter
 
 use datafusion::common::ScalarValue;
 use datafusion::logical_expr::ColumnarValue;
-use datafusion_functions_json::udfs::{json_contains_udf, json_get_str_udf};
+use datafusion_functions_json::udfs::{json_contains_udf, json_get_str_top_level_sorted_udf, json_get_str_udf};
 
 fn bench_json_contains(b: &mut Bencher) {
     let json_contains = json_contains_udf();
@@ -30,9 +30,54 @@ fn bench_json_get_str(b: &mut Bencher) {
     b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
 }
 
+fn make_json_negative_testcase() -> String {
+    // build a json with keys "b1", "b2" ... "b100", each with a large value ("a" repeated 1024 times)
+    let kvs = (0..100)
+        .map(|i| format!(r#""b{}": "{}""#, i, "a".repeat(1024)))
+        .collect::<Vec<_>>()
+        .join(",");
+    format!(r#"{{ {} }}"#, kvs)
+}
+
+fn bench_json_get_str_negative(b: &mut Bencher) {
+    let json_get_str = json_get_str_udf();
+    let args = &[
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some(make_json_negative_testcase()))),
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), // lexicographically less than "b1"
+    ];
+
+    b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
+}
+
+fn bench_json_get_str_sorted(b: &mut Bencher) {
+    let json_get_str = json_get_str_top_level_sorted_udf();
+    let args = &[
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+            r#"{"a": {"aa": "x", "ab: "y"}, "b": []}"#.to_string(),
+        ))),
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))),
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))),
+    ];
+
+    b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
+}
+
+fn bench_json_get_str_sorted_negative(b: &mut Bencher) {
+    let json_get_str = json_get_str_top_level_sorted_udf();
+    let args = &[
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some(make_json_negative_testcase()))),
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), // lexicographically less than "b1"
+    ];
+
+    b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("json_contains", bench_json_contains);
     c.bench_function("json_get_str", bench_json_get_str);
+    c.bench_function("json_get_str_negative", bench_json_get_str_negative);
+    c.bench_function("json_get_str_sorted", bench_json_get_str_sorted);
+    c.bench_function("json_get_str_sorted_negative", bench_json_get_str_sorted_negative);
 }
 
 criterion_group!(benches, criterion_benchmark);

src/common.rs

@@ -511,7 +511,11 @@ fn wrap_as_large_dictionary(original: &dyn AnyDictionaryArray, new_values: Array
     DictionaryArray::new(keys, new_values)
 }
 
-pub fn jiter_json_find<'j>(opt_json: Option<&'j str>, path: &[JsonPath]) -> Option<(Jiter<'j>, Peek)> {
+pub fn jiter_json_find<'j>(
+    opt_json: Option<&'j str>,
+    path: &[JsonPath],
+    mut sorted: Sortedness,
+) -> Option<(Jiter<'j>, Peek)> {
     let json_str = opt_json?;
     let mut jiter = Jiter::new(json_str.as_bytes());
     let mut peek = jiter.peek().ok()?;
@@ -521,6 +525,11 @@ pub fn jiter_json_find<'j>(opt_json: Option<&'j str>, path: &[JsonPath]) -> Opti
                 let mut next_key = jiter.known_object().ok()??;
 
                 while next_key != *key {
+                    if next_key > *key && matches!(sorted, Sortedness::Recursive | Sortedness::TopLevel) {
+                        // The current object is sorted and next_key is lexicographically greater than key
+                        // we are looking for, so we can early stop here.
+                        return None;
+                    }
                     jiter.next_skip().ok()?;
                     next_key = jiter.next_key().ok()??;
                 }
@@ -541,6 +550,9 @@ pub fn jiter_json_find<'j>(opt_json: Option<&'j str>, path: &[JsonPath]) -> Opti
                 return None;
             }
         }
+        if sorted == Sortedness::TopLevel {
+            sorted = Sortedness::Unspecified;
+        }
     }
     Some((jiter, peek))
 }
@@ -585,3 +597,27 @@ fn mask_dictionary_keys(keys: &PrimitiveArray<Int64Type>, type_ids: &[i8]) -> Pr
     }
     PrimitiveArray::new(keys.values().clone(), Some(null_mask.into()))
 }
+
+/// Information about the sortedness of a JSON object.
+/// This is used to optimize key lookups by early stopping when the key we are looking for is
+/// lexicographically greater than the current key and the object is known to be sorted.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub(crate) enum Sortedness {
+    /// No guarantees about the order of the elements.
+    Unspecified,
+    /// Only the outermost object is known to be sorted.
+    /// If the outermost item is not an object, this is equivalent to `Unspecified`.
+    TopLevel,
+    /// All objects are known to be sorted, including objects nested within arrays.
+    Recursive,
+}
+
+impl Sortedness {
+    pub(crate) fn function_name_suffix(self) -> &'static str {
+        match self {
+            Sortedness::Unspecified => "",
+            Sortedness::TopLevel => "_top_level_sorted",
+            Sortedness::Recursive => "_recursive_sorted",
+        }
+    }
+}

src/common_macros.rs

@@ -15,7 +15,7 @@
 ///
 /// [`ScalarUDFImpl`]: datafusion_expr::ScalarUDFImpl
 macro_rules! make_udf_function {
-    ($udf_impl:ty, $expr_fn_name:ident, $($arg:ident)*, $doc:expr) => {
+    ($udf_impl:ty, $expr_fn_name:ident, $($arg:ident)*, $doc:expr, $sorted:expr) => {
         paste::paste! {
             #[doc = $doc]
             #[must_use] pub fn $expr_fn_name($($arg: datafusion::logical_expr::Expr),*) -> datafusion::logical_expr::Expr {
@@ -37,7 +37,7 @@ macro_rules! make_udf_function {
                 [< STATIC_ $expr_fn_name:upper >]
                     .get_or_init(|| {
                         std::sync::Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl(
-                            <$udf_impl>::default(),
+                            <$udf_impl>::new($sorted),
                         ))
                     })
                     .clone()

src/json_as_text.rs

@@ -7,27 +7,48 @@ 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, return_type_check, GetError, InvokeResult, JsonPath};
+use crate::common::{
+    get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath, Sortedness,
+};
 use crate::common_macros::make_udf_function;
 
 make_udf_function!(
     JsonAsText,
     json_as_text,
     json_data path,
-    r#"Get any value from a JSON string by its "path", represented as a string"#
+    r#"Get any value from a JSON string by its "path", represented as a string"#,
+    Sortedness::Unspecified
+);
+
+make_udf_function!(
+    JsonAsText,
+    json_as_text_top_level_sorted,
+    json_data path,
+    r#"Get any value from a JSON string by its "path", represented as a string; assumes the JSON string's top level object's keys are sorted."#,
+    Sortedness::TopLevel
+);
+
+make_udf_function!(
+    JsonAsText,
+    json_as_text_recursive_sorted,
+    json_data path,
+    r#"Get any value from a JSON string by its "path", represented as a string; assumes all json object's keys are sorted."#,
+    Sortedness::Recursive
 );
 
 #[derive(Debug)]
 pub(super) struct JsonAsText {
     signature: Signature,
     aliases: [String; 1],
+    sorted: Sortedness,
 }
 
-impl Default for JsonAsText {
-    fn default() -> Self {
+impl JsonAsText {
+    pub fn new(sorted: Sortedness) -> Self {
         Self {
             signature: Signature::variadic_any(Volatility::Immutable),
-            aliases: ["json_as_text".to_string()],
+            aliases: [format!("json_as_text{}", sorted.function_name_suffix())],
+            sorted,
         }
     }
 }
@@ -50,7 +71,7 @@ impl ScalarUDFImpl for JsonAsText {
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
-        invoke::<StringArray>(args, jiter_json_as_text)
+        invoke::<StringArray>(args, |json, path| jiter_json_as_text(json, path, self.sorted))
     }
 
     fn aliases(&self) -> &[String] {
@@ -82,8 +103,8 @@ impl InvokeResult for StringArray {
     }
 }
 
-fn jiter_json_as_text(opt_json: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
-    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
+fn jiter_json_as_text(opt_json: Option<&str>, path: &[JsonPath], sorted: Sortedness) -> Result<String, GetError> {
+    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path, sorted) {
         match peek {
             Peek::Null => {
                 jiter.known_null()?;

src/json_contains.rs

@@ -14,20 +14,39 @@ make_udf_function!(
     JsonContains,
     json_contains,
     json_data path,
-    r#"Does the key/index exist within the JSON value as the specified "path"?"#
+    r#"Does the key/index exist within the JSON value as the specified "path"?"#,
+    crate::common::Sortedness::Unspecified
+);
+
+make_udf_function!(
+    JsonContains,
+    json_contains_top_level_sorted,
+    json_data path,
+    r#"Does the key/index exist within the JSON value as the specified "path"; assumes the JSON string's top level object's keys are sorted?"#,
+    crate::common::Sortedness::TopLevel
+);
+
+make_udf_function!(
+    JsonContains,
+    json_contains_recursive_sorted,
+    json_data path,
+    r#"Does the key/index exist within the JSON value as the specified "path"; assumes all json object's keys are sorted?"#,
+    crate::common::Sortedness::Recursive
 );
 
 #[derive(Debug)]
 pub(super) struct JsonContains {
     signature: Signature,
     aliases: [String; 1],
+    sorted: crate::common::Sortedness,
 }
 
-impl Default for JsonContains {
-    fn default() -> Self {
+impl JsonContains {
+    pub fn new(sorted: crate::common::Sortedness) -> Self {
         Self {
             signature: Signature::variadic_any(Volatility::Immutable),
-            aliases: ["json_contains".to_string()],
+            aliases: [format!("json_contains{}", sorted.function_name_suffix())],
+            sorted,
         }
     }
 }
@@ -54,7 +73,7 @@ impl ScalarUDFImpl for JsonContains {
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
-        invoke::<BooleanArray>(args, jiter_json_contains)
+        invoke::<BooleanArray>(args, |json, path| jiter_json_contains(json, path, self.sorted))
     }
 
     fn aliases(&self) -> &[String] {
@@ -88,6 +107,10 @@ impl InvokeResult for BooleanArray {
 }
 
 #[allow(clippy::unnecessary_wraps)]
-fn jiter_json_contains(json_data: Option<&str>, path: &[JsonPath]) -> Result<bool, GetError> {
-    Ok(jiter_json_find(json_data, path).is_some())
+fn jiter_json_contains(
+    json_data: Option<&str>,
+    path: &[JsonPath],
+    sorted: crate::common::Sortedness,
+) -> Result<bool, GetError> {
+    Ok(jiter_json_find(json_data, path, sorted).is_some())
 }

src/json_get.rs

@@ -10,6 +10,7 @@ use datafusion::scalar::ScalarValue;
 use jiter::{Jiter, NumberAny, NumberInt, Peek};
 
 use crate::common::InvokeResult;
+use crate::common::Sortedness;
 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};
@@ -18,22 +19,39 @@ make_udf_function!(
     JsonGet,
     json_get,
     json_data path,
-    r#"Get a value from a JSON string by its "path""#
+    r#"Get a value from a JSON string by its "path""#,
+    Sortedness::Unspecified
 );
 
-// build_typed_get!(JsonGet, "json_get", Union, Float64Array, jiter_json_get_float);
+make_udf_function!(
+    JsonGet,
+    json_get_top_level_sorted,
+    json_data path,
+    r#"Get a value from a JSON string by its "path"; assumes the JSON string's top level object's keys are sorted."#,
+    Sortedness::TopLevel
+);
+
+make_udf_function!(
+    JsonGet,
+    json_get_recursive_sorted,
+    json_data path,
+    r#"Get a value from a JSON string by its "path"; assumes all object's keys are sorted."#,
+    Sortedness::Recursive
+);
 
 #[derive(Debug)]
 pub(super) struct JsonGet {
     signature: Signature,
     aliases: [String; 1],
+    sorted: Sortedness,
 }
 
-impl Default for JsonGet {
-    fn default() -> Self {
+impl JsonGet {
+    pub fn new(sorted: Sortedness) -> Self {
         Self {
             signature: Signature::variadic_any(Volatility::Immutable),
-            aliases: ["json_get".to_string()],
+            aliases: [format!("json_get{}", sorted.function_name_suffix())],
+            sorted,
         }
     }
 }
@@ -56,7 +74,7 @@ impl ScalarUDFImpl for JsonGet {
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
-        invoke::<JsonUnion>(args, jiter_json_get_union)
+        invoke::<JsonUnion>(args, |json, path| jiter_json_get_union(json, path, self.sorted))
     }
 
     fn aliases(&self) -> &[String] {
@@ -93,8 +111,12 @@ impl InvokeResult for JsonUnion {
     }
 }
 
-fn jiter_json_get_union(opt_json: Option<&str>, path: &[JsonPath]) -> Result<JsonUnionField, GetError> {
-    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
+fn jiter_json_get_union(
+    opt_json: Option<&str>,
+    path: &[JsonPath],
+    sorted: Sortedness,
+) -> Result<JsonUnionField, GetError> {
+    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path, sorted) {
         build_union(&mut jiter, peek)
     } else {
         get_err!()