Merge 'Jsonb extract' from Ihor Andrianov

Made a jsonb traversal by json path.
Changed some ordinary json functions to use jsonb under the hood, so now
behavior of our json module more like sqlite.
Found and fixed some bugs on the way.

Closes #1135

COMPAT.md

@@ -361,14 +361,14 @@ Modifiers:
 | Function                           | Status  | Comment                                                                                                                                      |
 | ---------------------------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------- |
 | json(json)                         | Partial |                                                                                                                                              |
-| jsonb(json)                        |         |                                                                                                                                              |
+| jsonb(json)                        | Yes     |                                                                                                                                              |
 | json_array(value1,value2,...)      | Yes     |                                                                                                                                              |
 | jsonb_array(value1,value2,...)     |         |                                                                                                                                              |
 | json_array_length(json)            | Yes     |                                                                                                                                              |
 | json_array_length(json,path)       | Yes     |                                                                                                                                              |
 | json_error_position(json)          | Yes     |                                                                                                                                              |
 | json_extract(json,path,...)        | Partial | Does not fully support unicode literal syntax and does not allow numbers > 2^127 - 1 (which SQLite truncates to i32), does not support BLOBs |
-| jsonb_extract(json,path,...)       |         |                                                                                                                                              |
+| jsonb_extract(json,path,...)       | Yes     |                                                                                                                                              |
 | json -> path                       | Yes     |                                                                                                                                              |
 | json ->> path                      | Yes     |                                                                                                                                              |
 | json_insert(json,path,value,...)   |         |                                                                                                                                              |

core/function.rs

@@ -77,6 +77,7 @@ pub enum JsonFunc {
     JsonArrowExtract,
     JsonArrowShiftExtract,
     JsonExtract,
+    JsonbExtract,
     JsonObject,
     JsonType,
     JsonErrorPosition,
@@ -99,6 +100,7 @@ impl Display for JsonFunc {
                 Self::Jsonb => "jsonb".to_string(),
                 Self::JsonArray => "json_array".to_string(),
                 Self::JsonExtract => "json_extract".to_string(),
+                Self::JsonbExtract => "jsonb_extract".to_string(),
                 Self::JsonArrayLength => "json_array_length".to_string(),
                 Self::JsonArrowExtract => "->".to_string(),
                 Self::JsonArrowShiftExtract => "->>".to_string(),
@@ -559,6 +561,8 @@ impl Func {
             #[cfg(feature = "json")]
             "json_extract" => Ok(Func::Json(JsonFunc::JsonExtract)),
             #[cfg(feature = "json")]
+            "jsonb_extract" => Ok(Func::Json(JsonFunc::JsonbExtract)),
+            #[cfg(feature = "json")]
             "json_object" => Ok(Func::Json(JsonFunc::JsonObject)),
             #[cfg(feature = "json")]
             "json_type" => Ok(Func::Json(JsonFunc::JsonType)),

core/json/json_operations.rs

@@ -197,7 +197,7 @@ mod tests {
     }
 
     fn create_json(s: &str) -> OwnedValue {
-        OwnedValue::Text(Text::json(s))
+        OwnedValue::Text(Text::json(s.to_string()))
     }
 
     #[test]

core/json/json_path.rs

@@ -328,7 +328,9 @@ fn finalize_path<'a>(
         PPState::InKey => {
             if key_start < path.len() {
                 let key = &path[key_start..];
-                if key.starts_with('"') & !key.ends_with('"') {
+                if key.starts_with('"') && !key.ends_with('"') && key.len() > 1
+                    || (key.starts_with('"') && key.ends_with('"') && key.len() == 1)
+                {
                     bail_parse_error!("Bad json path: {}", path)
                 }
                 path_components.push(PathElement::Key(Cow::Borrowed(key), false));

core/json/jsonb.rs

@@ -1,6 +1,8 @@
 use crate::{bail_parse_error, LimboError, Result};
 use std::{fmt::Write, str::from_utf8};
 
+use super::json_path::{JsonPath, PathElement};
+
 const SIZE_MARKER_8BIT: u8 = 12;
 const SIZE_MARKER_16BIT: u8 = 13;
 const SIZE_MARKER_32BIT: u8 = 14;
@@ -190,6 +192,24 @@ pub enum ElementType {
     RESERVED3 = 15,
 }
 
+impl From<ElementType> for String {
+    fn from(element_type: ElementType) -> String {
+        match element_type {
+            ElementType::ARRAY => "array".to_string(),
+            ElementType::OBJECT => "object".to_string(),
+            ElementType::NULL => "null".to_string(),
+            ElementType::TRUE => "true".to_string(),
+            ElementType::FALSE => "false".to_string(),
+            ElementType::FLOAT | ElementType::FLOAT5 => "real".to_string(),
+            ElementType::INT | ElementType::INT5 => "integer".to_string(),
+            ElementType::TEXT | ElementType::TEXT5 | ElementType::TEXTJ | ElementType::TEXTRAW => {
+                "text".to_string()
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
 impl TryFrom<u8> for ElementType {
     type Error = LimboError;
 
@@ -218,10 +238,10 @@ impl TryFrom<u8> for ElementType {
 
 type PayloadSize = usize;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct JsonbHeader(ElementType, PayloadSize);
 
-enum HeaderFormat {
+pub(crate) enum HeaderFormat {
     Inline([u8; 1]),    // Small payloads embedded directly in the header
     OneByte([u8; 2]),   // Medium payloads with 1-byte size field
     TwoBytes([u8; 3]),  // Large payloads with 2-byte size field
@@ -229,7 +249,7 @@ enum HeaderFormat {
 }
 
 impl HeaderFormat {
-    fn as_bytes(&self) -> &[u8] {
+    pub fn as_bytes(&self) -> &[u8] {
         match self {
             Self::Inline(bytes) => bytes,
             Self::OneByte(bytes) => bytes,
@@ -244,6 +264,10 @@ impl JsonbHeader {
         Self(element_type, payload_size)
     }
 
+    pub fn make_null() -> Self {
+        Self(ElementType::NULL, 0)
+    }
+
     fn from_slice(cursor: usize, slice: &[u8]) -> Result<(Self, usize)> {
         match slice.get(cursor) {
             Some(header_byte) => {
@@ -294,7 +318,7 @@ impl JsonbHeader {
         }
     }
 
-    fn into_bytes(self) -> HeaderFormat {
+    pub fn into_bytes(self) -> HeaderFormat {
         let (element_type, payload_size) = (self.0, self.1);
 
         match payload_size {
@@ -359,17 +383,36 @@ impl Jsonb {
         self.data.len()
     }
 
+    pub fn make_empty_array(size: usize) -> Self {
+        let mut jsonb = Self {
+            data: Vec::with_capacity(size),
+        };
+        jsonb
+            .write_element_header(0, ElementType::ARRAY, 0)
+            .unwrap();
+        jsonb
+    }
+
+    pub fn append_to_array_unsafe(&mut self, data: &[u8]) {
+        self.data.extend_from_slice(data);
+    }
+
+    pub fn finalize_array_unsafe(&mut self) -> Result<()> {
+        self.write_element_header(0, ElementType::ARRAY, self.len() - 1)?;
+        Ok(())
+    }
+
     fn read_header(&self, cursor: usize) -> Result<(JsonbHeader, usize)> {
         let (header, offset) = JsonbHeader::from_slice(cursor, &self.data)?;
 
         Ok((header, offset))
     }
 
-    pub fn is_valid(&self) -> Result<()> {
+    pub fn is_valid(&self) -> Result<ElementType> {
         match self.read_header(0) {
             Ok((header, offset)) => {
-                if let Some(_) = self.data.get(offset..offset + header.1) {
-                    Ok(())
+                if self.data.get(offset..offset + header.1).is_some() {
+                    Ok(header.0)
                 } else {
                     bail_parse_error!("malformed JSON")
                 }
@@ -402,7 +445,7 @@ impl Jsonb {
             | JsonbHeader(ElementType::TEXTRAW, len)
             | JsonbHeader(ElementType::TEXTJ, len)
             | JsonbHeader(ElementType::TEXT5, len) => {
-                self.serialize_string(string, cursor, len, &header.0)?
+                self.serialize_string(string, cursor, len, &header.0, true)?
             }
             JsonbHeader(ElementType::INT, len)
             | JsonbHeader(ElementType::INT5, len)
@@ -436,7 +479,7 @@ impl Jsonb {
                 | ElementType::TEXTJ
                 | ElementType::TEXT5 => {
                     current_cursor =
-                        self.serialize_string(string, current_cursor, len, &element_type)?;
+                        self.serialize_string(string, current_cursor, len, &element_type, true)?;
                 }
                 _ => bail_parse_error!("malformed JSON"),
             }
@@ -474,9 +517,13 @@ impl Jsonb {
         cursor: usize,
         len: usize,
         kind: &ElementType,
+        quote: bool,
     ) -> Result<usize> {
         let word_slice = &self.data[cursor..cursor + len];
-        string.push('"');
+        if quote {
+            string.push('"');
+        }
+
         match kind {
             // Can be serialized as is. Do not need escaping
             ElementType::TEXT => {
@@ -630,7 +677,10 @@ impl Jsonb {
                 unreachable!()
             }
         }
-        string.push('"');
+        if quote {
+            string.push('"');
+        }
+
         Ok(cursor + len)
     }
 
@@ -833,6 +883,9 @@ impl Jsonb {
                 b',' if !first => {
                     pos += 1; // consume ','
                     pos = skip_whitespace(input, pos);
+                    if input[pos] == b',' {
+                        bail_parse_error!("2 commas in a row are not allowed")
+                    }
                 }
                 _ => {
                     // Parse key (must be string)
@@ -847,8 +900,11 @@ impl Jsonb {
                     pos = skip_whitespace(input, pos);
 
                     // Parse value - can be any JSON value including another object
-                    pos = self.deserialize_value(input, pos, depth)?;
-
+                    pos = self.deserialize_value(input, pos, depth + 1)?;
+                    pos = skip_whitespace(input, pos);
+                    if pos < input.len() && !matches!(input[pos], b',' | b'}') {
+                        bail_parse_error!("Should be , or }}")
+                    }
                     first = false;
                 }
             }
@@ -938,6 +994,8 @@ impl Jsonb {
 
             if quoted && c == quote {
                 break; // End of string
+            } else if !quoted && (c == b'"' || c == b'\'') {
+                bail_parse_error!("Something gone wrong")
             } else if c == b'\\' {
                 // Handle escape sequences
                 if pos >= input.len() {
@@ -1044,6 +1102,7 @@ impl Jsonb {
                         pos += 2;
                         element_type = ElementType::TEXT5;
                     }
+
                     _ => {
                         bail_parse_error!("Invalid escape sequence: \\{}", esc as char);
                     }
@@ -1367,9 +1426,150 @@ impl Jsonb {
         Ok(result)
     }
 
+    pub fn from_raw_data(data: &[u8]) -> Self {
+        Self::new(data.len(), Some(data))
+    }
+
     pub fn data(self) -> Vec<u8> {
         self.data
     }
+
+    pub fn get_by_path(&self, path: &JsonPath) -> Result<(Jsonb, ElementType)> {
+        let mut pos = 0;
+        let mut string_buffer = String::with_capacity(1024);
+        for segment in path.elements.iter() {
+            pos = self.navigate_to_segment(segment, pos, &mut string_buffer)?;
+        }
+        let (header, skip_header) = self.read_header(pos)?;
+        let end = pos + skip_header + header.1;
+        Ok((Jsonb::from_raw_data(&self.data[pos..end]), header.0))
+    }
+
+    pub fn get_by_path_raw(&self, path: &JsonPath) -> Result<&[u8]> {
+        let mut pos = 0;
+        let mut string_buffer = String::with_capacity(1024);
+        for segment in path.elements.iter() {
+            pos = self.navigate_to_segment(segment, pos, &mut string_buffer)?;
+        }
+        let (header, skip_header) = self.read_header(pos)?;
+        let end = pos + skip_header + header.1;
+        Ok(&self.data[pos..end])
+    }
+
+    pub fn array_len(&self) -> Result<usize> {
+        let (header, header_skip) = self.read_header(0)?;
+        if header.0 != ElementType::ARRAY {
+            return Ok(0);
+        }
+
+        let mut count = 0;
+        let mut pos = header_skip;
+        while pos < header_skip + header.1 {
+            pos = self.skip_element(pos)?;
+            count += 1;
+        }
+
+        Ok(count)
+    }
+
+    fn navigate_to_segment(
+        &self,
+        segment: &PathElement,
+        pos: usize,
+        string_buffer: &mut String,
+    ) -> Result<usize> {
+        let (header, skip_header) = self.read_header(pos)?;
+        let (parent_type, parent_size) = (header.0, header.1);
+        let mut current_pos = pos + skip_header;
+
+        match segment {
+            PathElement::Root() => return Ok(0),
+            PathElement::Key(path_key, is_raw) => {
+                if parent_type != ElementType::OBJECT {
+                    bail_parse_error!("parent is not object")
+                };
+                while current_pos < pos + parent_size {
+                    let (key_header, skip_header) = self.read_header(current_pos)?;
+                    let (key_type, key_size) = (key_header.0, key_header.1);
+                    current_pos += skip_header;
+
+                    if matches!(
+                        key_header.0,
+                        ElementType::TEXT
+                            | ElementType::TEXT5
+                            | ElementType::TEXTJ
+                            | ElementType::TEXTRAW
+                    ) {
+                        string_buffer.clear();
+                        current_pos = self.serialize_string(
+                            string_buffer,
+                            current_pos,
+                            key_size,
+                            &key_type,
+                            false,
+                        )?;
+
+                        if compare((&string_buffer, key_type), (path_key, *is_raw)) {
+                            return Ok(current_pos);
+                        } else {
+                            current_pos = self.skip_element(current_pos)?;
+                        }
+                    } else {
+                        bail_parse_error!("Key is not text!")
+                    };
+                }
+            }
+            PathElement::ArrayLocator(idx) => {
+                if parent_type != ElementType::ARRAY {
+                    bail_parse_error!("parent is not array");
+                };
+                if let Some(id) = idx {
+                    let id = id.to_owned();
+
+                    if id >= 0 {
+                        for _ in 0..id as usize {
+                            if current_pos < pos + parent_size {
+                                current_pos = self.skip_element(current_pos)?;
+                            } else {
+                                bail_parse_error!("Index is bigger then array size");
+                            }
+                        }
+                        return Ok(current_pos);
+                        // fix this after we remove serialized json
+                    } else {
+                        let mut temp_pos = current_pos;
+                        let mut count_elements = 0;
+                        while temp_pos < pos + parent_size {
+                            temp_pos = self.skip_element(temp_pos)?;
+                            count_elements += 1;
+                        }
+                        let corrected_idx = count_elements + id;
+                        if corrected_idx < 0 {
+                            bail_parse_error!("Index is bigger then array size")
+                        }
+                        for _ in 0..corrected_idx as usize {
+                            if current_pos < pos + parent_size {
+                                current_pos = self.skip_element(current_pos)?;
+                            } else {
+                                bail_parse_error!("Index is bigger then array size");
+                            }
+                        }
+                        return Ok(current_pos);
+                    }
+                } else {
+                    return Ok(pos);
+                }
+            }
+        }
+
+        bail_parse_error!("Not found")
+    }
+
+    fn skip_element(&self, mut pos: usize) -> Result<usize> {
+        let (header, skip_header) = self.read_header(pos)?;
+        pos += skip_header + header.1;
+        Ok(pos)
+    }
 }
 
 impl std::str::FromStr for Jsonb {
@@ -1380,6 +1580,140 @@ impl std::str::FromStr for Jsonb {
     }
 }
 
+#[inline]
+fn compare(key: (&str, ElementType), path_key: (&str, bool)) -> bool {
+    let (key, element_type) = key;
+    let (path_key, is_raw) = path_key;
+    if !is_raw && element_type == ElementType::TEXT {
+        if key.len() == path_key.len() {
+            return key == path_key;
+        } else {
+            return false;
+        }
+    }
+    if !is_raw {
+        return unescape_string(key) == path_key;
+    }
+    match element_type {
+        ElementType::TEXTJ | ElementType::TEXT5 | ElementType::TEXTRAW | ElementType::TEXT => {
+            return unescape_string(key) == unescape_string(path_key);
+        }
+        _ => {}
+    };
+
+    false
+}
+
+#[inline]
+pub fn unescape_string(input: &str) -> String {
+    let mut result = String::with_capacity(input.len());
+    let mut chars = input.chars().peekable();
+    let mut code_point = String::with_capacity(5);
+
+    while let Some(c) = chars.next() {
+        if c == '\\' {
+            match chars.next() {
+                Some('n') => result.push('\n'),
+                Some('r') => result.push('\r'),
+                Some('t') => result.push('\t'),
+                Some('\\') => result.push('\\'),
+                Some('/') => result.push('/'),
+                Some('"') => result.push('"'),
+                Some('b') => result.push('\u{0008}'),
+                Some('f') => result.push('\u{000C}'),
+                Some('x') => {
+                    code_point.clear();
+                    for _ in 0..2 {
+                        if let Some(hex_char) = chars.next() {
+                            code_point.push(hex_char);
+                        } else {
+                            break;
+                        }
+                    }
+                    if let Ok(code) = u16::from_str_radix(&code_point, 16) {
+                        if let Some(ch) = char::from_u32(code as u32) {
+                            result.push(ch)
+                        }
+                    }
+                }
+                // Handle \uXXXX format (JSON style)
+                Some('u') => {
+                    code_point.clear();
+                    for _ in 0..4 {
+                        if let Some(hex_char) = chars.next() {
+                            code_point.push(hex_char);
+                        } else {
+                            break;
+                        }
+                    }
+
+                    if let Ok(code) = u16::from_str_radix(&code_point, 16) {
+                        // Check if this is a high surrogate
+                        if matches!(code, 0xD800..=0xDBFF) {
+                            if chars.next() == Some('\\') && chars.next() == Some('u') {
+                                code_point.clear();
+                                for _ in 0..4 {
+                                    if let Some(hex_char) = chars.next() {
+                                        code_point.push(hex_char);
+                                    } else {
+                                        break;
+                                    }
+                                }
+
+                                if let Ok(low_code) = u16::from_str_radix(&code_point, 16) {
+                                    if (0xDC00..=0xDFFF).contains(&low_code) {
+                                        let high_ten_bits = (code - 0xD800) as u32;
+                                        let low_ten_bits = (low_code - 0xDC00) as u32;
+                                        let code_point = (high_ten_bits << 10) | low_ten_bits;
+                                        let unicode_value = code_point + 0x10000;
+
+                                        if let Some(unicode_char) = char::from_u32(unicode_value) {
+                                            result.push(unicode_char);
+                                        }
+                                    } else {
+                                        // If low surrogate is invalid, just push both as separate chars
+                                        if let Some(c1) = char::from_u32(code as u32) {
+                                            result.push(c1);
+                                        }
+                                        if let Some(c2) = char::from_u32(low_code as u32) {
+                                            result.push(c2);
+                                        }
+                                    }
+                                }
+                            } else {
+                                // No low surrogate, just push the high surrogate as is
+                                if let Some(unicode_char) = char::from_u32(code as u32) {
+                                    result.push(unicode_char);
+                                }
+                            }
+                        } else {
+                            // Not a surrogate pair, just a regular Unicode character
+                            if let Some(unicode_char) = char::from_u32(code as u32) {
+                                result.push(unicode_char);
+                            }
+                        }
+                    }
+                }
+
+                Some(c) => {
+                    // For any other escape sequence we don't recognize,
+                    // just output the backslash and the character
+                    result.push('\\');
+                    result.push(c);
+                }
+                None => {
+                    // Handle trailing backslash
+                    result.push('\\');
+                }
+            }
+        } else {
+            result.push(c);
+        }
+    }
+
+    result
+}
+
 #[inline]
 pub fn skip_whitespace(input: &[u8], mut pos: usize) -> usize {
     let len = input.len();
@@ -1742,6 +2076,10 @@ world""#,
         assert!(Jsonb::from_str("}").is_err());
         assert!(Jsonb::from_str("]").is_err());
 
+        assert!(Jsonb::from_str(r#"{"a":"55,"b":72}"#).is_err());
+
+        assert!(Jsonb::from_str(r#"{"a":"55",,"b":72}"#).is_err());
+
         // Unclosed string
         assert!(Jsonb::from_str(r#"{"key":"value"#).is_err());
 

core/json/mod.rs

@@ -5,6 +5,7 @@ mod json_path;
 mod jsonb;
 mod ser;
 
+use crate::bail_constraint_error;
 pub use crate::json::de::from_str;
 use crate::json::error::Error as JsonError;
 pub use crate::json::json_operations::{json_patch, json_remove};
@@ -13,7 +14,7 @@ pub use crate::json::ser::to_string;
 use crate::types::{OwnedValue, Text, TextSubtype};
 use crate::{bail_parse_error, json::de::ordered_object};
 use indexmap::IndexMap;
-use jsonb::Jsonb;
+use jsonb::{ElementType, Jsonb, JsonbHeader};
 use ser::to_string_pretty;
 use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
@@ -49,7 +50,7 @@ pub fn get_json(json_value: &OwnedValue, indent: Option<&str>) -> crate::Result<
                 None => to_string(&json_val)?,
             };
 
-            Ok(OwnedValue::Text(Text::json(&json)))
+            Ok(OwnedValue::Text(Text::json(json)))
         }
         OwnedValue::Blob(b) => {
             let jsonbin = Jsonb::new(b.len(), Some(b));
@@ -67,25 +68,13 @@ pub fn get_json(json_value: &OwnedValue, indent: Option<&str>) -> crate::Result<
                 None => to_string(&json_val)?,
             };
 
-            Ok(OwnedValue::Text(Text::json(&json)))
+            Ok(OwnedValue::Text(Text::json(json)))
         }
     }
 }
 
 pub fn jsonb(json_value: &OwnedValue) -> crate::Result<OwnedValue> {
-    let jsonbin = match json_value {
-        OwnedValue::Null | OwnedValue::Integer(_) | OwnedValue::Float(_) | OwnedValue::Text(_) => {
-            Jsonb::from_str(&json_value.to_text().unwrap())
-        }
-        OwnedValue::Blob(blob) => {
-            let blob = Jsonb::new(blob.len(), Some(&blob));
-            blob.is_valid()?;
-            Ok(blob)
-        }
-        _ => {
-            unimplemented!()
-        }
-    };
+    let jsonbin = convert_dbtype_to_jsonb(json_value);
     match jsonbin {
         Ok(jsonbin) => Ok(OwnedValue::Blob(Rc::new(jsonbin.data()))),
         Err(_) => {
@@ -94,6 +83,26 @@ pub fn jsonb(json_value: &OwnedValue) -> crate::Result<OwnedValue> {
     }
 }
 
+fn convert_dbtype_to_jsonb(val: &OwnedValue) -> crate::Result<Jsonb> {
+    match val {
+        OwnedValue::Text(text) => Jsonb::from_str(text.as_str()),
+        OwnedValue::Blob(blob) => {
+            let json = Jsonb::from_raw_data(blob);
+            json.is_valid()?;
+            Ok(json)
+        }
+        OwnedValue::Record(_) | OwnedValue::Agg(_) => {
+            bail_constraint_error!("Wront number of arguments");
+        }
+        OwnedValue::Null => Jsonb::from_str("null"),
+        OwnedValue::Float(float) => {
+            let mut buff = ryu::Buffer::new();
+            Jsonb::from_str(buff.format(*float))
+        }
+        OwnedValue::Integer(int) => Jsonb::from_str(&int.to_string()),
+    }
+}
+
 fn get_json_value(json_value: &OwnedValue) -> crate::Result<Val> {
     match json_value {
         OwnedValue::Text(ref t) => match from_str::<Val>(t.as_str()) {
@@ -149,29 +158,31 @@ pub fn json_array(values: &[OwnedValue]) -> crate::Result<OwnedValue> {
     }
 
     s.push(']');
-    Ok(OwnedValue::Text(Text::json(&s)))
+    Ok(OwnedValue::Text(Text::json(s)))
 }
 
 pub fn json_array_length(
     json_value: &OwnedValue,
     json_path: Option<&OwnedValue>,
 ) -> crate::Result<OwnedValue> {
-    let json = get_json_value(json_value)?;
+    let json = convert_dbtype_to_jsonb(json_value)?;
 
-    let arr_val = if let Some(path) = json_path {
-        match json_extract_single(&json, path, true)? {
-            Some(val) => val,
-            None => return Ok(OwnedValue::Null),
-        }
-    } else {
-        &json
-    };
-
-    match arr_val {
-        Val::Array(val) => Ok(OwnedValue::Integer(val.len() as i64)),
-        Val::Null => Ok(OwnedValue::Null),
-        _ => Ok(OwnedValue::Integer(0)),
+    if json_path.is_none() {
+        let result = json.array_len()?;
+        return Ok(OwnedValue::Integer(result as i64));
     }
+
+    let path = json_path_from_owned_value(json_path.expect("We already checked none"), true)?;
+
+    if let Some(path) = path {
+        if let Ok(len) = json
+            .get_by_path(&path)
+            .and_then(|(json, _)| json.array_len())
+        {
+            return Ok(OwnedValue::Integer(len as i64));
+        }
+    }
+    Ok(OwnedValue::Null)
 }
 
 pub fn json_set(json: &OwnedValue, values: &[OwnedValue]) -> crate::Result<OwnedValue> {
@@ -222,13 +233,14 @@ pub fn json_arrow_extract(value: &OwnedValue, path: &OwnedValue) -> crate::Resul
         return Ok(OwnedValue::Null);
     }
 
-    let json = get_json_value(value)?;
-    let extracted = json_extract_single(&json, path, false)?;
-
-    if let Some(val) = extracted {
-        let json = to_string(val)?;
-
-        Ok(OwnedValue::Text(Text::json(&json)))
+    if let Some(path) = json_path_from_owned_value(path, false)? {
+        let json = convert_dbtype_to_jsonb(value)?;
+        let extracted = json.get_by_path(&path);
+        if let Ok((json, _)) = extracted {
+            Ok(OwnedValue::Text(Text::json(json.to_string()?)))
+        } else {
+            Ok(OwnedValue::Null)
+        }
     } else {
         Ok(OwnedValue::Null)
     }
@@ -243,11 +255,17 @@ pub fn json_arrow_shift_extract(
     if let OwnedValue::Null = value {
         return Ok(OwnedValue::Null);
     }
-
-    let json = get_json_value(value)?;
-    let extracted = json_extract_single(&json, path, false)?.unwrap_or(&Val::Null);
-
-    convert_json_to_db_type(extracted, true)
+    if let Some(path) = json_path_from_owned_value(path, false)? {
+        let json = convert_dbtype_to_jsonb(value)?;
+        let extracted = json.get_by_path(&path);
+        if let Ok((json, element_type)) = extracted {
+            Ok(json_string_to_db_type(json, element_type, false)?)
+        } else {
+            Ok(OwnedValue::Null)
+        }
+    } else {
+        Ok(OwnedValue::Null)
+    }
 }
 
 /// Extracts a JSON value from a JSON object or array.
@@ -260,38 +278,99 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<O
 
     if paths.is_empty() {
         return Ok(OwnedValue::Null);
-    } else if paths.len() == 1 {
-        let json = get_json_value(value)?;
-        let extracted = json_extract_single(&json, &paths[0], true)?.unwrap_or(&Val::Null);
-
-        return convert_json_to_db_type(extracted, false);
     }
 
-    let json = get_json_value(value)?;
-    let mut result = "[".to_string();
+    let (json, element_type) = jsonb_extract_internal(value, paths)?;
+    let result = json_string_to_db_type(json, element_type, false)?;
 
-    for path in paths {
-        match path {
-            OwnedValue::Null => {
-                return Ok(OwnedValue::Null);
-            }
-            _ => {
-                let extracted = json_extract_single(&json, path, true)?.unwrap_or(&Val::Null);
-
-                if paths.len() == 1 && extracted == &Val::Null {
-                    return Ok(OwnedValue::Null);
-                }
-
-                result.push_str(&to_string(&extracted)?);
-                result.push(',');
+    Ok(result)
+}
+
+pub fn jsonb_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<OwnedValue> {
+    if let OwnedValue::Null = value {
+        return Ok(OwnedValue::Null);
+    }
+
+    if paths.is_empty() {
+        return Ok(OwnedValue::Null);
+    }
+
+    let (json, element_type) = jsonb_extract_internal(value, paths)?;
+    let result = json_string_to_db_type(json, element_type, true)?;
+
+    Ok(result)
+}
+
+fn jsonb_extract_internal(
+    value: &OwnedValue,
+    paths: &[OwnedValue],
+) -> crate::Result<(Jsonb, ElementType)> {
+    let null = Jsonb::from_raw_data(JsonbHeader::make_null().into_bytes().as_bytes());
+    if paths.len() == 1 {
+        if let Some(path) = json_path_from_owned_value(&paths[0], true)? {
+            let json = convert_dbtype_to_jsonb(value)?;
+            if let Ok((json, value_type)) = json.get_by_path(&path) {
+                return Ok((json, value_type));
+            } else {
+                return Ok((null, ElementType::NULL));
             }
+        } else {
+            return Ok((null, ElementType::NULL));
         }
     }
 
-    result.pop(); // remove the final comma
-    result.push(']');
+    let json = convert_dbtype_to_jsonb(value)?;
+    let mut result = Jsonb::make_empty_array(json.len());
 
-    Ok(OwnedValue::Text(Text::json(&result)))
+    let paths = paths
+        .into_iter()
+        .map(|p| json_path_from_owned_value(p, true));
+    for path in paths {
+        if let Some(path) = path? {
+            let fragment = json.get_by_path_raw(&path);
+            if let Ok(data) = fragment {
+                result.append_to_array_unsafe(data);
+            } else {
+                result.append_to_array_unsafe(JsonbHeader::make_null().into_bytes().as_bytes());
+            }
+        } else {
+            return Ok((null, ElementType::NULL));
+        }
+    }
+    result.finalize_array_unsafe()?;
+    Ok((result, ElementType::ARRAY))
+}
+
+fn json_string_to_db_type(
+    json: Jsonb,
+    element_type: ElementType,
+    raw_flag: bool,
+) -> crate::Result<OwnedValue> {
+    let mut json_string = json.to_string()?;
+    if raw_flag && matches!(element_type, ElementType::ARRAY | ElementType::OBJECT) {
+        return Ok(OwnedValue::Blob(Rc::new(json.data())));
+    }
+    match element_type {
+        ElementType::ARRAY | ElementType::OBJECT => Ok(OwnedValue::Text(Text::json(json_string))),
+        ElementType::TEXT | ElementType::TEXT5 | ElementType::TEXTJ | ElementType::TEXTRAW => {
+            json_string.remove(json_string.len() - 1);
+            json_string.remove(0);
+            Ok(OwnedValue::Text(Text {
+                value: Rc::new(json_string.into_bytes()),
+                subtype: TextSubtype::Text,
+            }))
+        }
+        ElementType::FLOAT5 | ElementType::FLOAT => Ok(OwnedValue::Float(
+            json_string.parse().expect("Should be valid f64"),
+        )),
+        ElementType::INT | ElementType::INT5 => Ok(OwnedValue::Integer(
+            json_string.parse().expect("Should be valid i64"),
+        )),
+        ElementType::TRUE => Ok(OwnedValue::Integer(1)),
+        ElementType::FALSE => Ok(OwnedValue::Integer(0)),
+        ElementType::NULL => Ok(OwnedValue::Null),
+        _ => unreachable!(),
+    }
 }
 
 /// Returns a value with type defined by SQLite documentation:
@@ -324,7 +403,7 @@ fn convert_json_to_db_type(extracted: &Val, all_as_db: bool) -> crate::Result<Ow
             if all_as_db {
                 Ok(OwnedValue::build_text(&json))
             } else {
-                Ok(OwnedValue::Text(Text::json(&json)))
+                Ok(OwnedValue::Text(Text::json(json)))
             }
         }
     }
@@ -357,91 +436,23 @@ pub fn json_type(value: &OwnedValue, path: Option<&OwnedValue>) -> crate::Result
     if let OwnedValue::Null = value {
         return Ok(OwnedValue::Null);
     }
+    if path.is_none() {
+        let json = convert_dbtype_to_jsonb(value)?;
+        let element_type = json.is_valid()?;
 
-    let json = get_json_value(value)?;
+        return Ok(OwnedValue::Text(Text::json(element_type.into())));
+    }
+    if let Some(path) = json_path_from_owned_value(path.unwrap(), true)? {
+        let json = convert_dbtype_to_jsonb(value)?;
 
-    let json = if let Some(path) = path {
-        match json_extract_single(&json, path, true)? {
-            Some(val) => val,
-            None => return Ok(OwnedValue::Null),
+        if let Ok((_, element_type)) = json.get_by_path(&path) {
+            Ok(OwnedValue::Text(Text::json(element_type.into())))
+        } else {
+            Ok(OwnedValue::Null)
         }
     } else {
-        &json
-    };
-
-    let val = match json {
-        Val::Null => "null",
-        Val::Bool(v) => {
-            if *v {
-                "true"
-            } else {
-                "false"
-            }
-        }
-        Val::Integer(_) => "integer",
-        Val::Float(_) => "real",
-        Val::String(_) => "text",
-        Val::Array(_) => "array",
-        Val::Object(_) => "object",
-        Val::Removed => unreachable!(),
-    };
-
-    Ok(OwnedValue::Text(Text::json(val)))
-}
-
-/// Returns the value at the given JSON path. If the path does not exist, it returns None.
-/// If the path is an invalid path, returns an error.
-///
-/// *strict* - if false, we will try to resolve the path even if it does not start with "$"
-///   in a way that's compatible with the `->` and `->>` operators. See examples in the docs:
-///   https://sqlite.org/json1.html#the_and_operators
-fn json_extract_single<'a>(
-    json: &'a Val,
-    path: &OwnedValue,
-    strict: bool,
-) -> crate::Result<Option<&'a Val>> {
-    let json_path = match json_path_from_owned_value(path, strict)? {
-        Some(path) => path,
-        None => return Ok(None),
-    };
-
-    let mut current_element = &Val::Null;
-
-    for element in json_path.elements.iter() {
-        match element {
-            PathElement::Root() => {
-                current_element = json;
-            }
-            PathElement::Key(key, _) => match current_element {
-                Val::Object(map) => {
-                    if let Some((_, value)) = map.iter().find(|(k, _)| k == key) {
-                        current_element = value;
-                    } else {
-                        return Ok(None);
-                    }
-                }
-                _ => return Ok(None),
-            },
-            PathElement::ArrayLocator(idx) => match current_element {
-                Val::Array(array) => {
-                    if let Some(mut idx) = *idx {
-                        if idx < 0 {
-                            idx += array.len() as i32;
-                        }
-
-                        if idx < array.len() as i32 {
-                            current_element = &array[idx as usize];
-                        } else {
-                            return Ok(None);
-                        }
-                    }
-                }
-                _ => return Ok(None),
-            },
-        }
+        Ok(OwnedValue::Null)
     }
-
-    Ok(Some(current_element))
 }
 
 fn json_path_from_owned_value(path: &OwnedValue, strict: bool) -> crate::Result<Option<JsonPath>> {
@@ -674,21 +685,16 @@ pub fn json_object(values: &[OwnedValue]) -> crate::Result<OwnedValue> {
         .collect::<Result<IndexMap<String, Val>, _>>()?;
 
     let result = crate::json::to_string(&value_map)?;
-    Ok(OwnedValue::Text(Text::json(&result)))
+    Ok(OwnedValue::Text(Text::json(result)))
 }
 
-pub fn is_json_valid(json_value: &OwnedValue) -> crate::Result<OwnedValue> {
-    match json_value {
-        OwnedValue::Text(ref t) => match from_str::<Val>(t.as_str()) {
-            Ok(_) => Ok(OwnedValue::Integer(1)),
-            Err(_) => Ok(OwnedValue::Integer(0)),
-        },
-        OwnedValue::Blob(_) => {
-            bail_parse_error!("Unsuported!")
-        }
-        OwnedValue::Null => Ok(OwnedValue::Null),
-        _ => Ok(OwnedValue::Integer(1)),
+pub fn is_json_valid(json_value: &OwnedValue) -> OwnedValue {
+    if matches!(json_value, OwnedValue::Null) {
+        return OwnedValue::Null;
     }
+    convert_dbtype_to_jsonb(json_value)
+        .map(|_| OwnedValue::Integer(1))
+        .unwrap_or(OwnedValue::Integer(0))
 }
 
 pub fn json_quote(value: &OwnedValue) -> crate::Result<OwnedValue> {
@@ -866,7 +872,7 @@ mod tests {
     #[test]
     fn test_json_array_simple() {
         let text = OwnedValue::build_text("value1");
-        let json = OwnedValue::Text(Text::json("\"value2\""));
+        let json = OwnedValue::Text(Text::json("\"value2\"".to_string()));
         let input = vec![text, json, OwnedValue::Integer(1), OwnedValue::Float(1.1)];
 
         let result = json_array(&input).unwrap();
@@ -1104,7 +1110,7 @@ mod tests {
         let text_key = OwnedValue::build_text("text_key");
         let text_value = OwnedValue::build_text("text_value");
         let json_key = OwnedValue::build_text("json_key");
-        let json_value = OwnedValue::Text(Text::json(r#"{"json":"value","number":1}"#));
+        let json_value = OwnedValue::Text(Text::json(r#"{"json":"value","number":1}"#.to_string()));
         let integer_key = OwnedValue::build_text("integer_key");
         let integer_value = OwnedValue::Integer(1);
         let float_key = OwnedValue::build_text("float_key");
@@ -1138,7 +1144,7 @@ mod tests {
     #[test]
     fn test_json_object_json_value_is_rendered_as_json() {
         let key = OwnedValue::build_text("key");
-        let value = OwnedValue::Text(Text::json(r#"{"json":"value"}"#));
+        let value = OwnedValue::Text(Text::json(r#"{"json":"value"}"#.to_string()));
         let input = vec![key, value];
 
         let result = json_object(&input).unwrap();

core/translate/expr.rs

@@ -894,16 +894,17 @@ pub fn translate_expr(
                             func_ctx,
                         )
                     }
-                    JsonFunc::JsonArray | JsonFunc::JsonExtract | JsonFunc::JsonSet => {
-                        translate_function(
-                            program,
-                            args.as_deref().unwrap_or_default(),
-                            referenced_tables,
-                            resolver,
-                            target_register,
-                            func_ctx,
-                        )
-                    }
+                    JsonFunc::JsonArray
+                    | JsonFunc::JsonExtract
+                    | JsonFunc::JsonSet
+                    | JsonFunc::JsonbExtract => translate_function(
+                        program,
+                        args.as_deref().unwrap_or_default(),
+                        referenced_tables,
+                        resolver,
+                        target_register,
+                        func_ctx,
+                    ),
                     JsonFunc::JsonArrowExtract | JsonFunc::JsonArrowShiftExtract => {
                         unreachable!(
                             "These two functions are only reachable via the -> and ->> operators"

core/types.rs

@@ -45,9 +45,9 @@ impl Text {
         }
     }
 
-    pub fn json(value: &str) -> Self {
+    pub fn json(value: String) -> Self {
         Self {
-            value: Rc::new(value.as_bytes().to_vec()),
+            value: Rc::new(value.into_bytes()),
             subtype: TextSubtype::Json,
         }
     }
@@ -186,7 +186,7 @@ impl OwnedValue {
                     return Ok(OwnedValue::Null);
                 };
                 if v.is_json() {
-                    Ok(OwnedValue::Text(Text::json(text)))
+                    Ok(OwnedValue::Text(Text::json(text.to_string())))
                 } else {
                     Ok(OwnedValue::build_text(text))
                 }

core/vdbe/mod.rs

@@ -53,6 +53,7 @@ use crate::{
     json::json_array_length, json::json_arrow_extract, json::json_arrow_shift_extract,
     json::json_error_position, json::json_extract, json::json_object, json::json_patch,
     json::json_quote, json::json_remove, json::json_set, json::json_type, json::jsonb,
+    json::jsonb_extract,
 };
 use crate::{info, CheckpointStatus};
 use crate::{
@@ -2174,6 +2175,23 @@ impl Program {
                                     Err(e) => return Err(e),
                                 }
                             }
+                            JsonFunc::JsonbExtract => {
+                                let result = match arg_count {
+                                    0 => jsonb_extract(&OwnedValue::Null, &[]),
+                                    _ => {
+                                        let val = &state.registers[*start_reg];
+                                        let reg_values = &state.registers
+                                            [*start_reg + 1..*start_reg + arg_count];
+
+                                        jsonb_extract(val, reg_values)
+                                    }
+                                };
+
+                                match result {
+                                    Ok(json) => state.registers[*dest] = json,
+                                    Err(e) => return Err(e),
+                                }
+                            }
                             JsonFunc::JsonArrowExtract | JsonFunc::JsonArrowShiftExtract => {
                                 assert_eq!(arg_count, 2);
                                 let json = &state.registers[*start_reg];
@@ -2218,7 +2236,7 @@ impl Program {
                             }
                             JsonFunc::JsonValid => {
                                 let json_value = &state.registers[*start_reg];
-                                state.registers[*dest] = is_json_valid(json_value)?;
+                                state.registers[*dest] = is_json_valid(json_value);
                             }
                             JsonFunc::JsonPatch => {
                                 assert_eq!(arg_count, 2);

testing/json.test

@@ -291,16 +291,15 @@ do_execsql_test json_extract_object_3 {
     SELECT json_extract('{"a": [1,2,3]}', '$.a', '$.a[0]', '$.a[1]', null, '$.a[3]')
 } {{}}
 
-# TODO: fix me - this passes on SQLite and needs to be fixed in Limbo.
-#       \x61 is the ASCII code for 'a'
-# do_execsql_test json_extract_with_escaping {
-#     SELECT json_extract('{"\x61": 1}', '$.a')
-# } {{1}}
 
-# TODO: fix me - this passes on SQLite and needs to be fixed in Limbo.
-#do_execsql_test json_extract_with_escaping_2 {
-#    SELECT json_extract('{"\x61": 1}', '$.\x61')
-#} {{1}}
+#       \x61 is the ASCII code for 'a'
+do_execsql_test json_extract_with_escaping {
+     SELECT json_extract('{"\x61": 1}', '$.a')
+} {{1}}
+
+do_execsql_test json_extract_with_escaping_2 {
+    SELECT json_extract('{"\x61": 1}', '$."\x61"')
+} {{1}}
 
 do_execsql_test json_extract_null_path {
     SELECT json_extract(1, null)
@@ -406,10 +405,10 @@ do_execsql_test json_arrow_shift_implicit_root_path {
   SELECT '{"a":1}' ->> 'a';
 } {{1}}
 
-# TODO: fix me after rebasing on top of https://github.com/tursodatabase/limbo/pull/631 - use the Option value in json_extract_single
-#do_execsql_test json_arrow_implicit_root_path_undefined_key {
-#  SELECT '{"a":1}' -> 'x';
-#} {{}}
+
+do_execsql_test json_arrow_implicit_root_path_undefined_key {
+  SELECT '{"a":1}' -> 'x';
+} {{}}
 
 do_execsql_test json_arrow_shift_implicit_root_path_undefined_key {
   SELECT '{"a":1}' ->> 'x';
@@ -476,10 +475,9 @@ do_execsql_test json_arrow_shift_array {
   SELECT '[1,2,3]' ->> '$'
 } {{[1,2,3]}}
 
-# TODO: fix me - this passes on SQLite and needs to be fixed in Limbo.
-#do_execsql_test json_extract_quote {
-#  SELECT json_extract('{"\"":1 }', '$.\"')
-#} {{1}}
+do_execsql_test json_extract_quote {
+  SELECT json_extract('{"\"":1 }', '$."\""')
+} {{1}}
 
 # Overflows 2**32 is equivalent to 0
 do_execsql_test json_extract_overflow_int32_1 {