implement json_array_length

COMPAT.md

@@ -234,8 +234,8 @@ Feature support of [sqlite expr syntax](https://www.sqlite.org/lang_expr.html).
 | jsonb(json)                        |         |         |
 | json_array(value1,value2,...)      | Yes     |         |
 | jsonb_array(value1,value2,...)     |         |         |
-| json_array_length(json)            |         |         |
-| json_array_length(json,path)       |         |         |
+| json_array_length(json)            | Yes     |         |
+| json_array_length(json,path)       | Yes     |         |
 | json_error_position(json)          |         |         |
 | json_extract(json,path,...)        |         |         |
 | jsonb_extract(json,path,...)       |         |         |

core/function.rs

@@ -6,6 +6,7 @@ use std::fmt::Display;
 pub enum JsonFunc {
     Json,
     JsonArray,
+    JsonArrayLength,
 }
 
 #[cfg(feature = "json")]
@@ -17,6 +18,7 @@ impl Display for JsonFunc {
             match self {
                 JsonFunc::Json => "json".to_string(),
                 JsonFunc::JsonArray => "json_array".to_string(),
+                JsonFunc::JsonArrayLength => "json_array_length".to_string(),
             }
         )
     }
@@ -334,6 +336,8 @@ impl Func {
             "json" => Ok(Func::Json(JsonFunc::Json)),
             #[cfg(feature = "json")]
             "json_array" => Ok(Func::Json(JsonFunc::JsonArray)),
+            #[cfg(feature = "json")]
+            "json_array_length" => Ok(Func::Json(JsonFunc::JsonArrayLength)),
             "unixepoch" => Ok(Func::Scalar(ScalarFunc::UnixEpoch)),
             "hex" => Ok(Func::Scalar(ScalarFunc::Hex)),
             "unhex" => Ok(Func::Scalar(ScalarFunc::Unhex)),

core/json/mod.rs

@@ -1,5 +1,6 @@
 mod de;
 mod error;
+mod path;
 mod ser;
 
 use std::rc::Rc;
@@ -8,9 +9,10 @@ pub use crate::json::de::from_str;
 pub use crate::json::ser::to_string;
 use crate::types::{LimboText, OwnedValue, TextSubtype};
 use indexmap::IndexMap;
+use path::get_json_val_by_path;
 use serde::{Deserialize, Serialize};
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
 #[serde(untagged)]
 pub enum Val {
     Null,
@@ -88,6 +90,49 @@ pub fn json_array(values: Vec<&OwnedValue>) -> crate::Result<OwnedValue> {
     Ok(OwnedValue::Text(LimboText::json(Rc::new(s))))
 }
 
+pub fn json_array_length(
+    json_value: &OwnedValue,
+    json_path: Option<&OwnedValue>,
+) -> crate::Result<OwnedValue> {
+    let path = match json_path {
+        Some(OwnedValue::Text(t)) => Some(t.value.to_string()),
+        Some(OwnedValue::Integer(i)) => Some(i.to_string()),
+        Some(OwnedValue::Float(f)) => Some(f.to_string()),
+        _ => None::<String>,
+    };
+
+    let top_val = match json_value {
+        OwnedValue::Text(ref t) => crate::json::from_str::<Val>(&t.value),
+        OwnedValue::Blob(b) => match jsonb::from_slice(b) {
+            Ok(j) => {
+                let json = j.to_string();
+                crate::json::from_str(&json)
+            }
+            Err(_) => crate::bail_parse_error!("malformed JSON"),
+        },
+        _ => return Ok(OwnedValue::Integer(0)),
+    };
+
+    let Ok(top_val) = top_val else {
+        crate::bail_parse_error!("malformed JSON")
+    };
+
+    let arr_val = if let Some(path) = path {
+        match get_json_val_by_path(&top_val, &path) {
+            Ok(Some(val)) => val,
+            Ok(None) => return Ok(OwnedValue::Null),
+            Err(e) => return Err(e),
+        }
+    } else {
+        &top_val
+    };
+
+    if let Val::Array(val) = &arr_val {
+        return Ok(OwnedValue::Integer(val.len() as i64));
+    }
+    Ok(OwnedValue::Integer(0))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -266,4 +311,121 @@ mod tests {
             Err(e) => assert!(e.to_string().contains("JSON cannot hold BLOB values")),
         }
     }
+
+    #[test]
+    fn test_json_array_length() {
+        let input = OwnedValue::build_text(Rc::new("[1,2,3,4]".to_string()));
+        let result = json_array_length(&input, None).unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 4);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_empty() {
+        let input = OwnedValue::build_text(Rc::new("[]".to_string()));
+        let result = json_array_length(&input, None).unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 0);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_root() {
+        let input = OwnedValue::build_text(Rc::new("[1,2,3,4]".to_string()));
+        let result = json_array_length(
+            &input,
+            Some(&OwnedValue::build_text(Rc::new("$".to_string()))),
+        )
+        .unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 4);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_not_array() {
+        let input = OwnedValue::build_text(Rc::new("{one: [1,2,3,4]}".to_string()));
+        let result = json_array_length(&input, None).unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 0);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_via_prop() {
+        let input = OwnedValue::build_text(Rc::new("{one: [1,2,3,4]}".to_string()));
+        let result = json_array_length(
+            &input,
+            Some(&OwnedValue::build_text(Rc::new("$.one".to_string()))),
+        )
+        .unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 4);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_via_index() {
+        let input = OwnedValue::build_text(Rc::new("[[1,2,3,4]]".to_string()));
+        let result = json_array_length(
+            &input,
+            Some(&OwnedValue::build_text(Rc::new("$[0]".to_string()))),
+        )
+        .unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 4);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_via_index_not_array() {
+        let input = OwnedValue::build_text(Rc::new("[1,2,3,4]".to_string()));
+        let result = json_array_length(
+            &input,
+            Some(&OwnedValue::build_text(Rc::new("$[2]".to_string()))),
+        )
+        .unwrap();
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 0);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
+
+    #[test]
+    fn test_json_array_length_via_index_bad_prop() {
+        let input = OwnedValue::build_text(Rc::new("{one: [1,2,3,4]}".to_string()));
+        let result = json_array_length(
+            &input,
+            Some(&OwnedValue::build_text(Rc::new("$.two".to_string()))),
+        )
+        .unwrap();
+        assert_eq!(OwnedValue::Null, result);
+    }
+
+    #[test]
+    fn test_json_array_length_simple_json_subtype() {
+        let input = OwnedValue::build_text(Rc::new("[1,2,3]".to_string()));
+        let wrapped = get_json(&input).unwrap();
+        let result = json_array_length(&wrapped, None).unwrap();
+
+        if let OwnedValue::Integer(res) = result {
+            assert_eq!(res, 3);
+        } else {
+            panic!("Expected OwnedValue::Integer");
+        }
+    }
 }

core/json/path.rs

@@ -0,0 +1,181 @@
+use super::Val;
+
+pub fn get_json_val_by_path<'v>(val: &'v Val, path: &str) -> crate::Result<Option<&'v Val>> {
+    match path.strip_prefix('$') {
+        Some(tail) => json_val_by_path(val, tail),
+        None => crate::bail_parse_error!("malformed path"),
+    }
+}
+
+fn json_val_by_path<'v>(val: &'v Val, path: &str) -> crate::Result<Option<&'v Val>> {
+    if path.is_empty() {
+        return Ok(Some(val));
+    }
+
+    match val {
+        Val::Array(inner) => {
+            if inner.is_empty() {
+                return Ok(None);
+            }
+            let Some(tail) = path.strip_prefix('[') else {
+                return Ok(None);
+            };
+            let (from_end, tail) = if let Some(updated_tail) = tail.strip_prefix("#-") {
+                (true, updated_tail)
+            } else {
+                (false, tail)
+            };
+
+            let Some((idx_str, tail)) = tail.split_once("]") else {
+                crate::bail_parse_error!("malformed path");
+            };
+
+            if idx_str.is_empty() {
+                return Ok(None);
+            }
+            let Ok(idx) = idx_str.parse::<usize>() else {
+                crate::bail_parse_error!("malformed path");
+            };
+            let result = if from_end {
+                inner.get(inner.len() - 1 - idx)
+            } else {
+                inner.get(idx)
+            };
+
+            if let Some(result) = result {
+                return json_val_by_path(result, tail);
+            }
+            Ok(None)
+        }
+        Val::Object(inner) => {
+            let Some(tail) = path.strip_prefix('.') else {
+                return Ok(None);
+            };
+
+            let (property, tail) = if let Some(tail) = tail.strip_prefix('"') {
+                if let Some((property, tail)) = tail.split_once('"') {
+                    (property, tail)
+                } else {
+                    crate::bail_parse_error!("malformed path");
+                }
+            } else if let Some(idx) = tail.find('.') {
+                (&tail[..idx], &tail[idx..])
+            } else {
+                (tail, "")
+            };
+
+            if let Some(result) = inner.get(property) {
+                return json_val_by_path(result, tail);
+            }
+            Ok(None)
+        }
+        _ => Ok(None),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_path_root() {
+        assert_eq!(
+            get_json_val_by_path(&Val::Bool(true), "$",).unwrap(),
+            Some(&Val::Bool(true))
+        );
+    }
+
+    #[test]
+    fn test_path_index() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Array(vec![Val::Integer(33), Val::Integer(55), Val::Integer(66)]),
+                "$[2]",
+            )
+            .unwrap(),
+            Some(&Val::Integer(66))
+        );
+    }
+
+    #[test]
+    fn test_path_negative_index() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Array(vec![Val::Integer(33), Val::Integer(55), Val::Integer(66)]),
+                "$[#-2]",
+            )
+            .unwrap(),
+            Some(&Val::Integer(33))
+        );
+    }
+
+    #[test]
+    fn test_path_index_deep() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Array(vec![Val::Array(vec![
+                    Val::Integer(33),
+                    Val::Integer(55),
+                    Val::Integer(66)
+                ])]),
+                "$[0][1]",
+            )
+            .unwrap(),
+            Some(&Val::Integer(55))
+        );
+    }
+
+    #[test]
+    fn test_path_prop_simple() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Object(
+                    [
+                        ("foo".into(), Val::Integer(55)),
+                        ("bar".into(), Val::Integer(66))
+                    ]
+                    .into()
+                ),
+                "$.bar",
+            )
+            .unwrap(),
+            Some(&Val::Integer(66))
+        );
+    }
+
+    #[test]
+    fn test_path_prop_nested() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Object(
+                    [(
+                        "foo".into(),
+                        Val::Object([("bar".into(), Val::Integer(66))].into())
+                    )]
+                    .into()
+                ),
+                "$.foo.bar",
+            )
+            .unwrap(),
+            Some(&Val::Integer(66))
+        );
+    }
+
+    #[test]
+    fn test_path_prop_quoted() {
+        assert_eq!(
+            get_json_val_by_path(
+                &Val::Object(
+                    [
+                        ("foo.baz".into(), Val::Integer(55)),
+                        ("bar".into(), Val::Integer(66))
+                    ]
+                    .into()
+                ),
+                r#"$."foo.baz""#,
+            )
+            .unwrap(),
+            Some(&Val::Integer(55))
+        );
+    }
+}

core/translate/expr.rs

@@ -913,6 +913,51 @@ pub fn translate_expr(
                         });
                         Ok(target_register)
                     }
+                    JsonFunc::JsonArrayLength => {
+                        let args = if let Some(args) = args {
+                            if args.len() > 2 {
+                                crate::bail_parse_error!(
+                                    "{} function with wrong number of arguments",
+                                    j.to_string()
+                                )
+                            }
+                            args
+                        } else {
+                            crate::bail_parse_error!(
+                                "{} function with no arguments",
+                                j.to_string()
+                            );
+                        };
+
+                        let json_reg = program.alloc_register();
+                        let path_reg = program.alloc_register();
+
+                        translate_expr(
+                            program,
+                            referenced_tables,
+                            &args[0],
+                            json_reg,
+                            precomputed_exprs_to_registers,
+                        )?;
+
+                        if args.len() == 2 {
+                            translate_expr(
+                                program,
+                                referenced_tables,
+                                &args[1],
+                                path_reg,
+                                precomputed_exprs_to_registers,
+                            )?;
+                        }
+
+                        program.emit_insn(Insn::Function {
+                            constant_mask: 0,
+                            start_reg: json_reg,
+                            dest: target_register,
+                            func: func_ctx,
+                        });
+                        Ok(target_register)
+                    }
                 },
                 Func::Scalar(srf) => {
                     match srf {

core/vdbe/mod.rs

@@ -37,7 +37,7 @@ use crate::types::{
 };
 use crate::util::parse_schema_rows;
 #[cfg(feature = "json")]
-use crate::{function::JsonFunc, json::get_json, json::json_array};
+use crate::{function::JsonFunc, json::get_json, json::json_array, json::json_array_length};
 use crate::{Connection, Result, TransactionState};
 use crate::{Rows, DATABASE_VERSION};
 use limbo_macros::Description;
@@ -2281,6 +2281,21 @@ impl Program {
                                 Err(e) => return Err(e),
                             }
                         }
+                        #[cfg(feature = "json")]
+                        crate::function::Func::Json(JsonFunc::JsonArrayLength) => {
+                            let json_value = &state.registers[*start_reg];
+                            let path_value = if arg_count > 1 {
+                                Some(&state.registers[*start_reg + 1])
+                            } else {
+                                None
+                            };
+                            let json_array_length = json_array_length(json_value, path_value);
+
+                            match json_array_length {
+                                Ok(length) => state.registers[*dest] = length,
+                                Err(e) => return Err(e),
+                            }
+                        }
                         crate::function::Func::Scalar(scalar_func) => match scalar_func {
                             ScalarFunc::Cast => {
                                 assert!(arg_count == 2);

testing/json.test

@@ -83,3 +83,35 @@ do_execsql_test json_array_json {
 do_execsql_test json_array_nested {
    SELECT json_array(json_array(1,2,3), json('[1,2,3]'), '[1,2,3]')
 } {{[[1,2,3],[1,2,3],"[1,2,3]"]}}
+
+do_execsql_test json_array_length {
+   SELECT json_array_length('[1,2,3,4]');
+} {{4}}
+
+do_execsql_test json_array_length_empty {
+   SELECT json_array_length('[]');
+} {{0}}
+
+do_execsql_test json_array_length_root {
+  SELECT json_array_length('[1,2,3,4]', '$');
+} {{4}}
+
+do_execsql_test json_array_length_not_array {
+  SELECT json_array_length('{"one":[1,2,3]}');
+} {{0}}
+
+do_execsql_test json_array_length_via_prop {
+  SELECT json_array_length('{"one":[1,2,3]}', '$.one');
+} {{3}}
+
+do_execsql_test json_array_length_via_index {
+  SELECT json_array_length('[[1,2,3,4]]', '$[0]');
+} {{4}}
+
+do_execsql_test json_array_length_via_index_not_array {
+  SELECT json_array_length('[1,2,3,4]', '$[2]');
+} {{0}}
+
+do_execsql_test json_array_length_via_bad_prop {
+  SELECT json_array_length('{"one":[1,2,3]}', '$.two');
+} {{}}